fix: updating point/surface export and constructor

adding properties, allowing file streams
adding from_dict and extending to_dict methods

Author: lachlangrose

LoopStructural/datatypes/_point.py

@@ -1,7 +1,8 @@
 from dataclasses import dataclass
 import numpy as np
 
-from typing import Optional
+from typing import Optional, Union
+import io
 
 
 @dataclass
@@ -16,6 +17,9 @@ def to_dict(self):
             "locations": self.locations,
             "values": self.values,
             "name": self.name,
+            "properties": (
+                {k: p.tolist() for k, p in self.properties.items()} if self.properties else None
+            ),
         }
 
     def vtk(self):
@@ -25,9 +29,14 @@ def vtk(self):
         points["values"] = self.values
         return points
 
-    def save(self, filename):
-        filename = str(filename)
-        ext = filename.split('.')[-1]
+    def save(self, filename: Union[str, io.StringIO], ext=None):
+        if isinstance(filename, io.StringIO):
+            if ext is None:
+                raise ValueError('Please provide an extension for StringIO')
+            ext = ext.lower()
+        else:
+            ext = filename.split('.')[-1].lower()
+            filename = str(filename)
         if ext == 'json':
             import json
 
@@ -49,9 +58,29 @@ def save(self, filename):
             from LoopStructural.export.gocad import _write_pointset
 
             _write_pointset(self, filename)
+        elif ext == 'csv':
+            import pandas as pd
+
+            df = pd.DataFrame(self.locations, columns=['x', 'y', 'z'])
+            df['value'] = self.values
+            if self.properties is not None:
+                for k, v in self.properties.items():
+                    df[k] = v
+            df.to_csv(filename, index=False)
         else:
             raise ValueError(f'Unknown file extension {ext}')
 
+    @classmethod
+    def from_dict(cls, d, flatten=False):
+        if 'locations' not in d:
+            raise ValueError('locations not in dictionary')
+        locations = np.array(d['locations'])
+        if flatten:
+            locations = locations.reshape((-1, 3))
+        return ValuePoints(
+            locations, d.get('values', None), d.get('name', 'unnamed'), d.get('properties', None)
+        )
+
 
 @dataclass
 class VectorPoints:
@@ -65,8 +94,14 @@ def to_dict(self):
             "locations": self.locations,
             "vectors": self.vectors,
             "name": self.name,
+            "properties": (
+                {k: p.tolist() for k, p in self.properties.items()} if self.properties else None
+            ),
         }
 
+    def from_dict(self, d):
+        return VectorPoints(d['locations'], d['vectors'], d['name'], d.get('properties', None))
+
     def vtk(self, geom='arrow', scale=1.0, scale_function=None, tolerance=0.05):
         import pyvista as pv
 
@@ -108,5 +143,16 @@ def save(self, filename):
             from LoopStructural.export.gocad import _write_pointset
 
             _write_pointset(self, filename)
+        elif ext == 'csv':
+            import pandas as pd
+
+            df = pd.DataFrame(self.locations, columns=['x', 'y', 'z'])
+            df['vx'] = self.vectors[:, 0]
+            df['vy'] = self.vectors[:, 1]
+            df['vz'] = self.vectors[:, 2]
+            if self.properties is not None:
+                for k, v in self.properties.items():
+                    df[k] = v
+            df.to_csv(filename)
         else:
             raise ValueError(f'Unknown file extension {ext}')

LoopStructural/datatypes/_surface.py

@@ -1,16 +1,18 @@
 from dataclasses import dataclass
 from typing import Optional
 import numpy as np
+import io
 
 
 @dataclass
 class Surface:
     vertices: np.ndarray
     triangles: np.ndarray
-    normals: np.ndarray
-    name: str
+    normals: Optional[np.ndarray] = None
+    name: str = 'surface'
     values: Optional[np.ndarray] = None
     properties: Optional[dict] = None
+    cell_properties: Optional[dict] = None
 
     @property
     def triangle_area(self):
@@ -75,26 +77,77 @@ def vtk(self):
         surface = pv.PolyData.from_regular_faces(self.vertices, self.triangles)
         if self.values is not None:
             surface["values"] = self.values
+        if self.properties is not None:
+            for k, v in self.properties.items():
+                surface.point_data[k] = np.array(v)
+        if self.cell_properties is not None:
+            for k, v in self.cell_properties.items():
+                surface.cell_data[k] = np.array(v)
         return surface
 
-    def to_dict(self):
+    def to_dict(self, flatten=False):
+        triangles = self.triangles
+        vertices = self.vertices
+        if flatten:
+            vertices = self.vertices.flatten()
+            triangles = (
+                np.hstack([np.ones((self.triangles.shape[0], 1)) * 3, self.triangles])
+                .astype(int)
+                .flatten()
+            )
         return {
-            "vertices": self.vertices,
-            "triangles": self.triangles,
-            "normals": self.normals,
+            "vertices": vertices.tolist(),
+            "triangles": triangles.tolist(),
+            "normals": self.normals.tolist() if self.normals is not None else None,
+            "properties": (
+                {k: p.tolist() for k, p in self.properties.items()} if self.properties else None
+            ),
+            "cell_properties": (
+                {k: p.tolist() for k, p in self.cell_properties.items()}
+                if self.cell_properties
+                else None
+            ),
             "name": self.name,
-            "values": self.values,
+            "values": self.values.tolist() if self.values is not None else None,
         }
 
-    def save(self, filename):
-        filename = str(filename)
-        ext = filename.split('.')[-1]
+    @classmethod
+    def from_dict(cls, d, flatten=False):
+        vertices = np.array(d['vertices'])
+        triangles = np.array(d['triangles'])
+        if flatten:
+            vertices = vertices.reshape((-1, 3))
+            triangles = triangles.reshape((-1, 4))[:, 1:]
+        return cls(
+            vertices,
+            triangles,
+            np.array(d['normals']),
+            d['name'],
+            np.array(d['values']),
+            d.get('properties', None),
+            d.get('cell_properties', None),
+        )
+
+    def save(self, filename, ext=None):
+        import pyvista as pv
+
+        print(filename, ext)
+        if isinstance(filename, (io.StringIO, io.BytesIO)):
+            if ext is None:
+                raise ValueError('Please provide an extension for StringIO')
+            ext = ext.lower()
+        else:
+            filename = str(filename)
+            if ext is None:
+                ext = filename.split('.')[-1].lower()
+
         if ext == 'json':
             import json
 
             with open(filename, 'w') as f:
                 json.dump(self.to_dict(), f)
         elif ext == 'vtk':
+            print(filename)
             self.vtk().save(filename)
         elif ext == 'obj':
             import meshio
@@ -105,7 +158,7 @@ def save(self, filename):
                 [("triangle", self.triangles)],
                 point_data={"normals": self.normals},
             )
-        elif ext == 'ts':
+        elif ext == 'ts' or ext == 'gocad':
             from LoopStructural.export.exporters import _write_feat_surfs_gocad
 
             _write_feat_surfs_gocad(self, filename)
@@ -119,5 +172,13 @@ def save(self, filename):
 
             with open(filename, 'wb') as f:
                 pickle.dump(self, f)
+        elif ext == 'csv':
+            import pandas as pd
+
+            df = pd.DataFrame(self.vertices, columns=['x', 'y', 'z'])
+            if self.properties:
+                for k, v in self.properties.items():
+                    df[k] = v
+            df.to_csv(filename, index=False)
         else:
             raise ValueError(f"Extension {ext} not supported")