Add mesh support (#7)

.github/workflows/test_and_deploy.yml

@@ -38,10 +38,11 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install setuptools pytest
+          pip install setuptools pytest pipx
           pip install -e .
+          pipx install hatch
       - name: Test with pytest
-        run: pytest -v -p no:warnings tests
+        run: hatch run all:test -v -p no:warnings tests
         env:
           PLATFORM: ${{ matrix.platform }}
 

README.md

@@ -12,6 +12,8 @@ in Python.
 
 ## Usage
 
+### As pandas DataFrame
+
 ```python
 import imodmodel
 
@@ -31,6 +33,31 @@ Out[3]:
 
 ```
 
+### As ImodModel object
+
+```python
+from imodmodel import ImodModel
+model = ImodModel.from_file("my_model_file.mod")
+```
+
+```ipython
+In [3]: model.objects[0].contours[0].points
+Out[3]: 
+array([[  6.875,  62.875, 124.   ], ...])
+
+In [4]: model.objects[0].meshes[0].vertices
+Out[4]: 
+array([[ 6.87500000e+00,  6.28750000e+01,  1.24000000e+02], ...])
+
+In [5]: model.objects[0].meshes[0].indices
+Out[5]: 
+array([[156,  18, 152], ...])
+
+In [6]: model.objects[0].meshes[0].face_values
+Out[6]: 
+array([0., 0., 35.22094345, ...])
+```
+
 That's it!
 
 ## Installation

pyproject.toml

@@ -100,6 +100,26 @@ extend-ignore = [
 [tool.ruff.per-file-ignores]
 "tests/*.py" = ["D"]
 
+[tool.hatch.envs.default]
+dependencies = [
+  "pytest",
+]
+
+[tool.hatch.envs.default.scripts]
+test = "pytest {args:tests}"
+
+[tool.hatch.envs.all]
+matrix-name-format = "{variable}_{value}"
+
+[[tool.hatch.envs.all.matrix]]
+pydantic_version = ["1","2"]
+
+[tool.hatch.envs.all.overrides]
+matrix.pydantic_version.dependencies = [
+    { value="pydantic<2", if = ["1"] },
+    { value="pydantic>=2", if = ["2"] }
+]
+
 # https://docs.pytest.org/en/6.2.x/customize.html
 [tool.pytest.ini_options]
 minversion = "6.0"

src/imodmodel/__init__.py

@@ -1,3 +1,4 @@
 from .functions import read
+from .models import ImodModel
 
-__all__ = ["read"]
+__all__ = ["read", "ImodModel"]

src/imodmodel/binary_specification.py

@@ -61,6 +61,13 @@ class ModFileSpecification:
         'time': 'i',
         'surf': 'i',
     }
+    MESH_HEADER = {
+        'vsize': 'i',
+        'lsize': 'i',
+        'flag': 'I',
+        'time': 'h',
+        'surf': 'h',
+    }
     IMAT = {
         'ambient': 'B',
         'diffuse': 'B',
@@ -104,10 +111,6 @@ class ModFileSpecification:
     OLBL = NotImplemented
     CLIP = NotImplemented
     MCLP = NotImplemented
-    MOST = NotImplemented
-    OBST = NotImplemented
-    COST = NotImplemented
-    MEST = NotImplemented
     SLAN = NotImplemented
     MEPA = NotImplemented
     SKLI = NotImplemented

src/imodmodel/models.py

@@ -1,8 +1,10 @@
 import os
-from typing import Tuple, List, Optional
+import warnings
+from typing import Tuple, List, Optional, Union
 
 import numpy as np
 from pydantic import BaseModel, validator
+from pydantic.version import VERSION as PYDANTIC_VERSION
 
 
 class ID(BaseModel):
@@ -11,6 +13,18 @@ class ID(BaseModel):
     version_id: str
 
 
+class GeneralStorage(BaseModel):
+    """https://bio3d.colorado.edu/imod/doc/binspec.html"""
+    type: int
+    flags: int
+    index: Union[float, int, Tuple[int, int], Tuple[int, int, int, int]]
+    value: Union[float, int, Tuple[int, int], Tuple[int, int, int, int]]
+
+    if PYDANTIC_VERSION < '2.0':
+        class Config:
+            smart_union = True
+
+
 class ModelHeader(BaseModel):
     """https://bio3d.colorado.edu/imod/doc/binspec.html"""
     name: str
@@ -88,10 +102,80 @@ class Contour(BaseModel):
     """https://bio3d.colorado.edu/imod/doc/binspec.html"""
     header: ContourHeader
     points: np.ndarray  # pt
+    extra: List[GeneralStorage] = []
+
+    class Config:
+        arbitrary_types_allowed = True
+
+
+class MeshHeader(BaseModel):
+    """https://bio3d.colorado.edu/imod/doc/binspec.html"""
+    vsize: int
+    lsize: int
+    flag: int
+    time: int
+    surf: int
+
+class Mesh(BaseModel):
+    """https://bio3d.colorado.edu/imod/doc/binspec.html"""
+    header: MeshHeader
+    raw_vertices: np.ndarray
+    raw_indices: np.ndarray
+    extra: List[GeneralStorage] = []
 
     class Config:
         arbitrary_types_allowed = True
 
+    @validator('raw_indices')
+    @classmethod
+    def validate_indices(cls, indices: np.ndarray):
+        if indices.ndim > 1:
+            raise ValueError('indices must be 1D')
+        if indices[-1] != -1:
+            raise ValueError('Indices must end with -1')
+        if len(indices[np.where(indices >= 0)]) % 3 != 0:
+            raise ValueError(f'Invalid indices shape: {indices.shape}')
+        for i in (-20, -23, -24):
+            if i in indices:
+                warnings.warn(f'Unsupported mesh type: {i}')
+        return indices
+
+    @validator('raw_vertices')
+    @classmethod
+    def validate_vertices(cls, vertices: np.ndarray):
+        if vertices.ndim > 1:
+            raise ValueError('vertices must be 1D')
+        if len(vertices) % 3 != 0:
+            raise ValueError(f'Invalid vertices shape: {vertices.shape}')
+        return vertices
+
+    @property
+    def vertices(self) -> np.ndarray:
+        return self.raw_vertices.reshape((-1, 3))
+
+    @property
+    def indices(self) -> np.ndarray:
+        return self.raw_indices[np.where(self.raw_indices >= 0)].reshape((-1, 3))
+
+    @property
+    def face_values(self) -> Optional[np.ndarray]:
+        """Extra value for each vertex face.
+        The extra values are index, value  pairs
+        However, the index is an index into the indices array,
+        not directly an index of a vertex.
+        Furthermore, the index has to be fixed because
+        the original indices array has special command values (-25, -22, -1, ...)
+        """
+        values = np.zeros((len(self.vertices),))
+        has_face_values = False
+        for extra in self.extra:
+            if not (extra.type == 10 and isinstance(extra.index, int)):
+                continue
+            has_face_values = True
+            values[self.raw_indices[extra.index]] = extra.value
+        if has_face_values:
+            return values
+
 
 class IMAT(BaseModel):
     """https://bio3d.colorado.edu/imod/doc/binspec.html"""
@@ -145,6 +229,8 @@ class View(BaseModel):
 class Object(BaseModel):
     """https://bio3d.colorado.edu/imod/doc/binspec.html"""
     contours: List[Contour] = []
+    meshes: List[Mesh] = []
+    extra: List[GeneralStorage] = []
 
 
 class ImodModel(BaseModel):
@@ -156,6 +242,7 @@ class ImodModel(BaseModel):
     header: ModelHeader
     objects: List[Object]
     imat: Optional[IMAT]
+    extra: List[GeneralStorage] = []
 
     @classmethod
     def from_file(cls, filename: os.PathLike):

src/imodmodel/parsers.py

@@ -1,6 +1,6 @@
 import re
 from struct import Struct
-from typing import Any, BinaryIO, Dict, Tuple
+from typing import Any, BinaryIO, Dict, List, Tuple, Union
 
 import numpy as np
 
@@ -9,6 +9,9 @@
     IMAT,
     Contour,
     ContourHeader,
+    GeneralStorage,
+    Mesh,
+    MeshHeader,
     ImodModel,
     ModelHeader,
     Object,
@@ -40,6 +43,28 @@ def _parse_from_format_str(file: BinaryIO, format_str: str) -> Tuple[Any, ...]:
     return struct.unpack(file.read(struct.size))
 
 
+def _parse_from_type_flags(file: BinaryIO, flags: int) -> Union[int, float, Tuple[int, int], Tuple[int, int, int, int]]:
+    """Determine the next type from a flag, and parse the correct type.
+    The general storage chunks (MOST, OBST, MEST, COST) carry values as type unions.
+    The type of the union is stored in a 2-bit flag:
+    - 0b00: int
+    - 0b01: float
+    - 0b10: short, short
+    - 0b11: byte, byte, byte, byte
+    """
+    flag_mask, flag_int, flag_float, flag_short, flag_byte = 0b11, 0b00, 0b01, 0b10, 0b11
+    if flags & flag_mask == flag_int:
+        return _parse_from_format_str(file, '>i')[0]
+    elif flags & flag_mask == flag_float:
+        return _parse_from_format_str(file, '>f')[0]
+    elif flags & flag_mask == flag_short:
+        return _parse_from_format_str(file, '>2h')
+    elif flags & flag_mask == flag_byte:
+        return _parse_from_format_str(file, '>4b')
+    else:
+        raise ValueError(f'Invalid flags: {flags}')
+
+
 def _parse_id(file: BinaryIO) -> ID:
     data = _parse_from_specification(file, ModFileSpecification.ID)
     return ID(**data)
@@ -56,12 +81,8 @@ def _parse_object_header(file: BinaryIO) -> ObjectHeader:
 
 
 def _parse_object(file: BinaryIO) -> Object:
-    header = _parse_object_header(file)
-    contours = []
-    for _ in range(header.contsize):
-        _parse_control_sequence(file)
-        contours.append(_parse_contour(file))
-    return Object(contours=contours)
+    _parse_object_header(file)
+    return Object()
 
 
 def _parse_contour_header(file: BinaryIO) -> ContourHeader:
@@ -76,6 +97,24 @@ def _parse_contour(file: BinaryIO) -> Contour:
     return Contour(header=header, points=pt)
 
 
+def _parse_mesh_header(file: BinaryIO) -> MeshHeader:
+    data = _parse_from_specification(file, ModFileSpecification.MESH_HEADER)
+    return MeshHeader(**data)
+
+
+def _parse_mesh(file: BinaryIO) -> Mesh:
+    header = _parse_mesh_header(file)
+    vertices = _parse_from_format_str(file, f">{'fff' * header.vsize}")
+    vertices = np.array(vertices)
+    indices = _parse_from_format_str(file, f">{'i' * header.lsize}")
+    indices = np.array(indices)
+    # Only support the simplest mesh case, which is that each polygon
+    # starts with -25 and ends with -22, and the list is terminated by -1
+    if any(i in indices for i in (-20, -21, -23, -24)):
+        raise ValueError("This mesh type is not yet supported")
+    return Mesh(header=header, raw_vertices=vertices, raw_indices=indices)
+
+
 def _parse_control_sequence(file: BinaryIO) -> str:
     return file.read(4).decode("utf-8")
 
@@ -91,6 +130,20 @@ def _parse_imat(file: BinaryIO) -> IMAT:
     return IMAT(**data)
 
 
+def _parse_general_storage(file: BinaryIO) -> List[GeneralStorage]:
+    size = _parse_chunk_size(file)
+    if size % 12 != 0:
+        raise ValueError(f"Chunk size not divisible by 12: {size}")
+    storages = list()
+    n_chunks = size // 12
+    for _ in range(n_chunks):
+        type, flags = _parse_from_format_str(file, '>hh')
+        index = _parse_from_type_flags(file, flags)
+        value = _parse_from_type_flags(file, flags>>2)
+        storages.append(GeneralStorage(type=type, flags=flags, index=index, value=value))
+    return storages
+
+
 def _parse_unknown(file: BinaryIO) -> None:
     bytes_to_skip = _parse_chunk_size(file)
     file.read(bytes_to_skip)
@@ -101,16 +154,27 @@ def parse_model(file: BinaryIO) -> ImodModel:
     header = _parse_model_header(file)
     control_sequence = _parse_control_sequence(file)
     imat = None
+    extra = list()
 
     objects = []
     while control_sequence != "IEOF":
         if control_sequence == "OBJT":
             objects.append(_parse_object(file))
         elif control_sequence == "IMAT":
             imat = _parse_imat(file)
+        elif control_sequence == "CONT":
+            objects[-1].contours.append(_parse_contour(file))
         elif control_sequence == "MESH":
-            break
+            objects[-1].meshes.append(_parse_mesh(file))
+        elif control_sequence == "MOST":
+            extra += _parse_general_storage(file)
+        elif control_sequence == "OBST":
+            objects[-1].extra += _parse_general_storage(file)
+        elif control_sequence == "COST":
+            objects[-1].contours[-1].extra += _parse_general_storage(file)
+        elif control_sequence == "MEST":
+            objects[-1].meshes[-1].extra += _parse_general_storage(file)
         else:
             _parse_unknown(file)
         control_sequence = _parse_control_sequence(file)
-    return ImodModel(id=id, header=header, objects=objects, imat=imat)
+    return ImodModel(id=id, header=header, objects=objects, imat=imat, extra=extra)

tests/conftest.py

@@ -20,3 +20,15 @@ def two_contour_model_file_handle(two_contour_model_file):
 def meshed_contour_model_file() -> Path:
     """A model file with meshed contours."""
     return TEST_DATA_DIRECTORY / 'meshed_contour_example.mod'
+
+
+@pytest.fixture
+def meshed_curvature_model_file() -> Path:
+    """A model file with curvature measurements from imodcurvature"""
+    return TEST_DATA_DIRECTORY / 'meshed_curvature_example.mod'
+
+
+@pytest.fixture
+def meshed_curvature_model_file_handle(meshed_curvature_model_file):
+    """A file handle with curvature measurements from imodcurvature"""
+    return open(meshed_curvature_model_file, mode='rb')