Merge pull request #139 from DCC-Lab/jlb/export

Add raw data export

Author: JLBegin

pytissueoptics/rayscattering/energyLogging/energyLogger.py

@@ -1,6 +1,7 @@
+import json
 import os
 import pickle
-from typing import List, Union, Optional
+from typing import List, Optional, TextIO, Union
 
 import numpy as np
 
@@ -10,12 +11,15 @@
 from pytissueoptics.rayscattering.scatteringScene import ScatteringScene
 from pytissueoptics.scene.geometry import Vector
 from pytissueoptics.scene.logger.listArrayContainer import ListArrayContainer
-from pytissueoptics.scene.logger.logger import DataType, InteractionKey, Logger
+from pytissueoptics.scene.logger.logger import DataType, InteractionData, InteractionKey, Logger
 
+from ..opencl.CLScene import NO_SOLID_LABEL
 from .energyType import EnergyType
 
 
 class EnergyLogger(Logger):
+    _data: dict[InteractionKey, InteractionData]
+
     def __init__(
         self,
         scene: ScatteringScene,
@@ -307,3 +311,78 @@ def _fluenceTransform(self, key: InteractionKey, data: Optional[np.ndarray]) ->
 
         data[:, 0] = data[:, 0] / self._scene.getMaterial(key.solidLabel).mu_a
         return data
+
+    def export(self, exportPath: str):
+        """
+        Export the raw 3D data points to a CSV file, along with the scene information to a JSON file.
+
+        The data file <exportPath>.csv will be comma-delimited and will contain the following columns:
+        - energy, x, y, z, solid_index, surface_index
+
+        Two types of interactions are logged: scattering and surface crossings. In the first case, the energy will be
+        the delta energy deposited at the point and the surface index will be -1. In the second case, the energy
+        will be the total photon energy when crossing the surface, either as positive if leaving the surface
+        (along the normal) or as negative if entering the surface.
+
+        The scene information will be saved in a JSON file named <exportPath>.json, which includes details for each solid
+        index and surface index, such as their labels, materials, and geometry. The world information is also exported
+        as solid index -1.
+        """
+        if not self.has3D:
+            utils.warn("Cannot export data when keep3D is False. No 3D data available.")
+            return
+
+        solidLabels = []
+        for solid in self._scene.solids:
+            if solid.isStack():
+                solidLabels.extend(solid.getLayerLabels())
+            else:
+                solidLabels.append(solid.getLabel())
+        solidLabels.sort()
+
+        print("Exporting raw data to file...")
+        filepath = f"{exportPath}.csv"
+        with open(filepath, "w") as file:
+            file.write("energy,x,y,z,solid_index,surface_index\n")
+            self._writeKeyData(file, InteractionKey(NO_SOLID_LABEL), -1, -1)
+            for i, solidLabel in enumerate(solidLabels):
+                self._writeKeyData(file, InteractionKey(solidLabel), i, -1)
+                for j, surfaceLabel in enumerate(self._scene.getSurfaceLabels(solidLabel)):
+                    self._writeKeyData(file, InteractionKey(solidLabel, surfaceLabel), i, j)
+        print(f"Exported data points to {filepath}")
+
+        sceneInfo = {}
+        material = self._scene.getWorldEnvironment().material
+        sceneInfo["-1"] = {"label": "world", "material": material.__dict__ if material else None}
+        for i, solidLabel in enumerate(solidLabels):
+            material = self._scene.getMaterial(solidLabel)
+            solid = self._scene.getSolid(solidLabel)
+            surfaces = {}
+            for j, surfaceLabel in enumerate(solid.surfaceLabels):
+                normals = [s.normal for s in solid.getPolygons(surfaceLabel)[:2]]
+                if len(normals) == 1 or normals[0] == normals[1]:
+                    normal = normals[0].array
+                else:
+                    normal = None
+                surfaces[j] = {"label": surfaceLabel, "normal": normal}
+
+            sceneInfo[str(i)] = {
+                "label": solidLabel,
+                "type": solid.__class__.__name__,
+                "material": material.__dict__ if material else None,
+                "geometry": solid.geometryExport(),
+                "surfaces": surfaces,
+            }
+
+        sceneFilepath = f"{exportPath}.json"
+        with open(sceneFilepath, "w") as file:
+            json.dump(sceneInfo, file, indent=4)
+        print(f"Exported scene information to {sceneFilepath}")
+
+    def _writeKeyData(self, file: TextIO, key: InteractionKey, solidIndex: int, surfaceIndex: int):
+        if key not in self._data or self._data[key].dataPoints is None:
+            return
+        dataArray = self._data[key].dataPoints.getData().astype(str)
+        dataArray = np.hstack((dataArray, np.full((dataArray.shape[0], 2), str(solidIndex))))
+        dataArray[:, 5] = str(surfaceIndex)
+        file.write("\n".join([",".join(row) for row in dataArray]) + "\n")

pytissueoptics/rayscattering/tests/energyLogging/testEnergyLogger.py

@@ -1,11 +1,13 @@
 import io
+import json
 import os
 import tempfile
 import unittest
 from unittest.mock import MagicMock, patch
 
 import numpy as np
 
+from pytissueoptics import Sphere
 from pytissueoptics.rayscattering.display.utils import Direction
 from pytissueoptics.rayscattering.display.views import (
     View2DProjection,
@@ -16,6 +18,8 @@
 )
 from pytissueoptics.rayscattering.energyLogging import EnergyLogger
 from pytissueoptics.rayscattering.materials import ScatteringMaterial
+from pytissueoptics.rayscattering.opencl.CLScene import NO_SOLID_LABEL
+from pytissueoptics.rayscattering.samples import PhantomTissue
 from pytissueoptics.rayscattering.scatteringScene import ScatteringScene
 from pytissueoptics.scene.geometry import Vector
 from pytissueoptics.scene.logger import InteractionKey
@@ -284,3 +288,64 @@ def testWhenLogSegmentArray_shouldRaiseError(self):
     def testWhenLogSegment_shouldRaiseError(self):
         with self.assertRaises(NotImplementedError):
             self.logger.logSegment(Vector(0, 0, 0), Vector(0, 0, 0), self.INTERACTION_KEY)
+
+    def testWhenExport_shouldExport3DDataPointsToFile(self):
+        # Use a scene that contains a stack, a sphere and a world material.
+        scene = PhantomTissue(worldMaterial=ScatteringMaterial(0.1, 0.1, 0.99))
+        scene.add(Sphere(position=Vector(0, 5, 0), material=ScatteringMaterial(0.4, 0.2, 0.9)))
+        self.logger = EnergyLogger(scene)
+
+        # Log entering surface event, world scattering event and scattering event in both solids.
+        self.logger.logDataPoint(0.1, Vector(0.7, 0.8, 0.8), InteractionKey("middleLayer"))
+        self.logger.logDataPoint(-0.9, Vector(0.5, 1.0, 0.75), InteractionKey("frontLayer", "interface1"))
+        self.logger.logDataPoint(0.4, Vector(0, 5, 0), InteractionKey("sphere"))
+        self.logger.logDataPoint(0.2, Vector(0, 0, 0), InteractionKey(NO_SOLID_LABEL))
+
+        with tempfile.TemporaryDirectory() as tempDir:
+            filePath = os.path.join(tempDir, "test_sim")
+            self.logger.export(filePath)
+            self.assertTrue(os.path.exists(filePath + ".csv"))
+
+            with open(filePath + ".csv", "r") as f:
+                lines = f.readlines()
+
+            self.assertEqual(5, len(lines))
+            self.assertEqual("energy,x,y,z,solid_index,surface_index\n", lines[0])
+            self.assertEqual("0.2,0.0,0.0,0.0,-1,-1\n", lines[1])
+            self.assertEqual("-0.9,0.5,1.0,0.75,1,5\n", lines[2])
+            self.assertEqual("0.1,0.7,0.8,0.8,2,-1\n", lines[3])
+            self.assertEqual("0.4,0.0,5.0,0.0,3,-1\n", lines[4])
+
+    def testWhenExport_shouldExportMetadataToFile(self):
+        scene = PhantomTissue(worldMaterial=ScatteringMaterial(0.1, 0.1, 0.99))
+        scene.add(Sphere(position=Vector(0, 5, 0), material=ScatteringMaterial(0.4, 0.2, 0.9)))
+        self.logger = EnergyLogger(scene)
+
+        with tempfile.TemporaryDirectory() as tempDir:
+            filePath = os.path.join(tempDir, "test_sim")
+            self.logger.export(filePath)
+            self.assertTrue(os.path.exists(filePath + ".json"))
+            sceneInfo = json.loads(open(filePath + ".json", "r").read())
+
+        self.assertEqual(["-1", "0", "1", "2", "3"], list(sceneInfo.keys()))
+
+        expectedWorldInfo = {
+            "label": "world",
+            "material": {
+                "mu_s": 0.1,
+                "mu_a": 0.1,
+                "mu_t": 0.2,
+                "_albedo": 0.5,
+                "g": 0.99,
+                "n": 1.0,
+            },
+        }
+        self.assertEqual(expectedWorldInfo, sceneInfo["-1"])
+
+        self.assertEqual(["label", "type", "material", "geometry", "surfaces"], list(sceneInfo["0"].keys()))
+        self.assertEqual("backLayer", sceneInfo["0"]["label"])
+        self.assertEqual("Cuboid", sceneInfo["0"]["type"])
+        expectedLayerGeometry = {"shape": [3, 3, 2], "position": [0, 0, 1], "bbox": [[-1.5, 1.5], [-1.5, 1.5], [0, 2]]}
+        self.assertEqual(expectedLayerGeometry, sceneInfo["0"]["geometry"])
+        self.assertEqual(16, len(sceneInfo["0"]["surfaces"]))
+        self.assertEqual({"label": "interface0", "normal": [0.0, 0.0, -1.0]}, sceneInfo["0"]["surfaces"]["14"])

pytissueoptics/scene/solids/cuboid.py

@@ -151,3 +151,6 @@ def _getLayerPolygons(self, layerLabel: str) -> List[Polygon]:
         for surfaceLabel in layerSurfaceLabels:
             polygons.extend(self._surfaces.getPolygons(surfaceLabel))
         return polygons
+
+    def _geometryParams(self) -> dict:
+        return {"shape": self.shape}

pytissueoptics/scene/solids/cylinder.py

@@ -164,3 +164,12 @@ def __hash__(self):
         materialHash = hash(self._material) if self._material else 0
         propertyHash = hash((self._radius, self._length, self._frontCenter, self._backCenter))
         return hash((materialHash, propertyHash))
+
+    def _geometryParams(self) -> dict:
+        return {
+            "radius": self._radius,
+            "length": self._length,
+            "u": self._u,
+            "v": self._v,
+            "s": self._s,
+        }

pytissueoptics/scene/solids/ellipsoid.py

@@ -212,3 +212,11 @@ def _getRadiusError(self) -> float:
 
     def _getMinimumRadiusTowards(self, vertex) -> float:
         return (1 - self._getRadiusError()) * self._radiusTowards(vertex)
+
+    def _geometryParams(self) -> dict:
+        return {
+            "radius_a": self._a,
+            "radius_b": self._b,
+            "radius_c": self._c,
+            "order": self._order,
+        }

pytissueoptics/scene/solids/lens.py

@@ -50,6 +50,7 @@ def __init__(
         self._diameter = diameter
         self._frontRadius = frontRadius
         self._backRadius = backRadius
+        self._thickness = thickness
 
         length = self._computeEdgeThickness(thickness)
         super().__init__(
@@ -156,6 +157,18 @@ def smooth(self, surfaceLabel: str = None, reset: bool = True):
         for surfaceLabel in ["front", "back"]:
             self.smooth(surfaceLabel, reset=False)
 
+    def _geometryParams(self) -> dict:
+        return {
+            "diameter": self._diameter,
+            "frontRadius": self._frontRadius,
+            "backRadius": self._backRadius,
+            "thickness": self._thickness,
+            "length": self._length,
+            "u": self._u,
+            "v": self._v,
+            "s": self._s,
+        }
+
 
 class SymmetricLens(ThickLens):
     """A symmetrical thick lens of focal length `f` in air."""

pytissueoptics/scene/solids/solid.py

@@ -1,6 +1,6 @@
 import warnings
 from functools import partial
-from typing import Callable, Dict, List
+from typing import Any, Callable, Dict, List
 
 import numpy as np
 
@@ -332,3 +332,20 @@ def __hash__(self):
         verticesHash = hash(tuple(sorted([hash(v) for v in self._vertices])))
         materialHash = hash(self._material) if self._material else 0
         return hash((verticesHash, materialHash))
+
+    def geometryExport(self) -> dict[str, Any]:
+        """Used to describe geometry during data export."""
+        params = {
+            **self._geometryParams(),
+            "position": self.position.array,
+            "bbox": self.getBoundingBox().xyzLimits,
+        }
+        if self._orientation != INITIAL_SOLID_ORIENTATION:
+            params["orientation"] = self._orientation.array
+        if self._rotation:
+            params["rotation"] = [self._rotation.xTheta, self._rotation.yTheta, self._rotation.zTheta]
+        return params
+
+    def _geometryParams(self) -> dict:
+        """To be implemented by Solid subclasses to detail other geometry parameters."""
+        return {}

pytissueoptics/scene/solids/sphere.py

@@ -59,3 +59,9 @@ def _getMinimumRadius(self) -> float:
 
     def _radiusTowards(self, vertex) -> float:
         return self.radius
+
+    def _geometryParams(self) -> dict:
+        return {
+            "radius": self._radius,
+            "order": self._order,
+        }