Skip to content

Commit 3198ea0

Browse files
authored
Merge pull request #82 from Loop3D/copilot/add-obj-ascii-voxel-export
Add ASCII voxel export and fix surface export format detection
2 parents 960ac02 + 8b8adc8 commit 3198ea0

File tree

2 files changed

+212
-4
lines changed

2 files changed

+212
-4
lines changed

loopstructural/gui/visualisation/object_list_widget.py

Lines changed: 67 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
import logging
2+
13
import pyvista as pv
24
from PyQt5.QtCore import Qt
35
from PyQt5.QtWidgets import (
@@ -13,6 +15,8 @@
1315
QWidget,
1416
)
1517

18+
logger = logging.getLogger(__name__)
19+
1620

1721
class ObjectListWidget(QWidget):
1822
def __init__(self, parent=None, *, viewer=None, properties_widget=None):
@@ -262,11 +266,19 @@ def export_selected_object(self):
262266
except ImportError:
263267
has_geoh5py = False
264268

265-
if hasattr(object, "faces"): # Likely a surface/mesh
269+
# Check if this is a grid/voxel type (UniformGrid, ImageData, StructuredGrid, RectilinearGrid)
270+
is_grid = type(mesh).__name__ in ['UniformGrid', 'ImageData', 'StructuredGrid', 'RectilinearGrid']
271+
272+
if is_grid:
273+
# Grid/voxel meshes support ASCII export
274+
formats = ["vtk", "ascii"]
275+
if has_geoh5py:
276+
formats.append("geoh5")
277+
elif hasattr(mesh, "faces"): # Likely a surface/mesh
266278
formats = ["obj", "vtk", "ply"]
267279
if has_geoh5py:
268280
formats.append("geoh5")
269-
elif hasattr(object, "points"): # Likely a point cloud
281+
elif hasattr(mesh, "points"): # Likely a point cloud
270282
formats = ["vtp"]
271283
if has_geoh5py:
272284
formats.append("geoh5")
@@ -279,6 +291,7 @@ def export_selected_object(self):
279291
"vtk": "VTK (*.vtk)",
280292
"ply": "PLY (*.ply)",
281293
"vtp": "VTP (*.vtp)",
294+
"ascii": "ASCII Grid (*.txt)",
282295
"geoh5": "Geoh5 (*.geoh5)",
283296
}
284297
filters = ";;".join([filter_map[f] for f in formats])
@@ -312,6 +325,9 @@ def export_selected_object(self):
312325
if hasattr(mesh, "save")
313326
else pv.save_meshio(file_path, mesh)
314327
)
328+
elif selected_format == "ascii":
329+
# Export grid/voxel as ASCII: x, y, z, value format
330+
self._export_grid_ascii(mesh, file_path, object_label)
315331
elif selected_format == "geoh5":
316332
with geoh5py.Geoh5(file_path, overwrite=True) as geoh5:
317333
if hasattr(mesh, "faces"):
@@ -320,11 +336,58 @@ def export_selected_object(self):
320336
)
321337
else:
322338
geoh5.add_points(name=object_label, vertices=mesh.points)
323-
print(f"Exported {object_label} to {file_path} as {selected_format}")
339+
logger.info(f"Exported {object_label} to {file_path} as {selected_format}")
324340
except Exception as e:
325-
print(f"Failed to export object: {e}")
341+
logger.error(f"Failed to export object: {e}")
326342
# Logic for exporting the object
327343

344+
def _export_grid_ascii(self, mesh, file_path, object_label):
345+
"""Export a grid/voxel mesh to ASCII format.
346+
347+
Format: x, y, z, value (one line per cell center)
348+
349+
Parameters
350+
----------
351+
mesh : pyvista grid mesh
352+
The grid mesh to export
353+
file_path : str
354+
Path to the output file
355+
object_label : str
356+
Name of the object (used to determine which scalar array to export)
357+
"""
358+
import numpy as np
359+
360+
# Get cell centers
361+
cell_centers = mesh.cell_centers()
362+
centers = cell_centers.points
363+
364+
# Get scalar values - try to use the active scalars or the first available array
365+
scalar_name = mesh.active_scalars_name
366+
if scalar_name is None:
367+
# Try to find any cell data array
368+
if mesh.cell_data:
369+
scalar_name = list(mesh.cell_data.keys())[0]
370+
371+
if scalar_name is not None:
372+
values = mesh.cell_data[scalar_name]
373+
else:
374+
# If no scalar data, use zeros
375+
values = np.zeros(mesh.n_cells)
376+
377+
# Write to file
378+
with open(file_path, 'w') as f:
379+
f.write(f"# ASCII Grid Export: {object_label}\n")
380+
f.write(f"# Format: x y z value\n")
381+
f.write(f"# Number of cells: {mesh.n_cells}\n")
382+
if scalar_name:
383+
f.write(f"# Scalar field: {scalar_name}\n")
384+
f.write("#\n")
385+
386+
for i in range(len(centers)):
387+
x, y, z = centers[i]
388+
value = values[i]
389+
f.write(f"{x:.6f} {y:.6f} {z:.6f} {value:.6f}\n")
390+
328391
def remove_selected_object(self):
329392
selected_items = self.treeWidget.selectedItems()
330393
if not selected_items:

tests/unit/test_grid_export.py

Lines changed: 145 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,145 @@
1+
"""Unit tests for grid export functionality."""
2+
3+
import tempfile
4+
from pathlib import Path
5+
from unittest.mock import MagicMock
6+
7+
import numpy as np
8+
import pytest
9+
10+
11+
def test_export_grid_ascii_basic():
12+
"""Test ASCII export of a simple grid with scalar data."""
13+
# Create a mock mesh object that simulates a PyVista UniformGrid
14+
mock_mesh = MagicMock()
15+
mock_mesh.n_cells = 8
16+
17+
# Mock cell centers
18+
mock_cell_centers = MagicMock()
19+
mock_cell_centers.points = np.array([
20+
[0.5, 0.5, 0.5],
21+
[1.5, 0.5, 0.5],
22+
[0.5, 1.5, 0.5],
23+
[1.5, 1.5, 0.5],
24+
[0.5, 0.5, 1.5],
25+
[1.5, 0.5, 1.5],
26+
[0.5, 1.5, 1.5],
27+
[1.5, 1.5, 1.5],
28+
])
29+
mock_mesh.cell_centers.return_value = mock_cell_centers
30+
31+
# Mock scalar data
32+
mock_mesh.active_scalars_name = "scalar_field"
33+
mock_mesh.cell_data = {
34+
"scalar_field": np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0])
35+
}
36+
37+
# Create a temporary file for export
38+
with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False) as f:
39+
temp_path = f.name
40+
41+
try:
42+
# Import the ObjectListWidget to get the export method
43+
# We'll test the export method in isolation
44+
from loopstructural.gui.visualisation.object_list_widget import ObjectListWidget
45+
46+
# Create a minimal instance (viewer and properties_widget can be None for this test)
47+
widget = ObjectListWidget(viewer=MagicMock(), properties_widget=None)
48+
49+
# Call the export method
50+
widget._export_grid_ascii(mock_mesh, temp_path, "test_grid")
51+
52+
# Read the exported file
53+
with open(temp_path, 'r') as f:
54+
lines = f.readlines()
55+
56+
# Verify header
57+
assert lines[0].strip() == "# ASCII Grid Export: test_grid"
58+
assert lines[1].strip() == "# Format: x y z value"
59+
assert lines[2].strip() == "# Number of cells: 8"
60+
assert lines[3].strip() == "# Scalar field: scalar_field"
61+
assert lines[4].strip() == "#"
62+
63+
# Verify data lines
64+
data_lines = lines[5:]
65+
assert len(data_lines) == 8
66+
67+
# Verify first data line
68+
first_line = data_lines[0].strip().split()
69+
assert len(first_line) == 4
70+
assert float(first_line[0]) == pytest.approx(0.5, abs=1e-5)
71+
assert float(first_line[1]) == pytest.approx(0.5, abs=1e-5)
72+
assert float(first_line[2]) == pytest.approx(0.5, abs=1e-5)
73+
assert float(first_line[3]) == pytest.approx(1.0, abs=1e-5)
74+
75+
finally:
76+
# Clean up
77+
Path(temp_path).unlink(missing_ok=True)
78+
79+
80+
def test_export_grid_ascii_no_scalars():
81+
"""Test ASCII export when grid has no scalar data."""
82+
# Create a mock mesh without scalar data
83+
mock_mesh = MagicMock()
84+
mock_mesh.n_cells = 4
85+
86+
mock_cell_centers = MagicMock()
87+
mock_cell_centers.points = np.array([
88+
[0.5, 0.5, 0.5],
89+
[1.5, 0.5, 0.5],
90+
[0.5, 1.5, 0.5],
91+
[1.5, 1.5, 0.5],
92+
])
93+
mock_mesh.cell_centers.return_value = mock_cell_centers
94+
95+
# No scalar data
96+
mock_mesh.active_scalars_name = None
97+
mock_mesh.cell_data = {}
98+
99+
with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False) as f:
100+
temp_path = f.name
101+
102+
try:
103+
from loopstructural.gui.visualisation.object_list_widget import ObjectListWidget
104+
105+
widget = ObjectListWidget(viewer=MagicMock(), properties_widget=None)
106+
widget._export_grid_ascii(mock_mesh, temp_path, "test_grid_no_scalars")
107+
108+
with open(temp_path, 'r') as f:
109+
lines = f.readlines()
110+
111+
# Should still create file with zeros
112+
assert "# ASCII Grid Export: test_grid_no_scalars" in lines[0]
113+
assert len(lines) >= 9 # Header + 4 data lines
114+
115+
# Verify that values are zero
116+
data_lines = lines[5:]
117+
for line in data_lines:
118+
if line.strip():
119+
parts = line.strip().split()
120+
if len(parts) == 4:
121+
value = float(parts[3])
122+
assert value == pytest.approx(0.0, abs=1e-5)
123+
124+
finally:
125+
Path(temp_path).unlink(missing_ok=True)
126+
127+
128+
def test_mesh_type_detection():
129+
"""Test that grid mesh types are correctly detected."""
130+
test_cases = [
131+
("UniformGrid", True),
132+
("ImageData", True),
133+
("StructuredGrid", True),
134+
("RectilinearGrid", True),
135+
("PolyData", False),
136+
("UnstructuredGrid", False),
137+
]
138+
139+
for mesh_type, expected_is_grid in test_cases:
140+
mock_mesh = MagicMock()
141+
mock_mesh.__class__.__name__ = mesh_type
142+
143+
is_grid = type(mock_mesh).__name__ in ['UniformGrid', 'ImageData', 'StructuredGrid', 'RectilinearGrid']
144+
145+
assert is_grid == expected_is_grid, f"Failed for mesh type: {mesh_type}"

0 commit comments

Comments
 (0)