test_extraction.py

Go to the documentation of this file.

"""Smoke test for pore-network extraction from an SDF.
 
Run:  PYTHONPATH=<pnm/build> python scripts/test_extraction.py <sdf.vti>
(e.g. ../flow/data/packing_ring.vti — 7199 pores). Writes the pores to pores.vtp.
"""
import sys
 
import peclet.pnm as pnm
 
def test_extraction(filename):
    print(f"Reading {filename}...")
    sdf_3d, origin_zyx, spacing_zyx = pnm.SDFReader.read_vti(filename)
    print(f"Loaded SDF. Shape: {sdf_3d.shape}")
 
    print("Extracting pores on GPU...")
    pores = pnm.extract_pores(sdf_3d, origin_zyx, spacing_zyx)
    
    print(f"Found {len(pores)} pores.")
    if len(pores) > 0:
        print("First 5 pores:")
        for i in range(min(5, len(pores))):
            print(pores[i])
            
    # Simple validation: Check if duplicate locations?
    import numpy as np
    coords = np.array([[p.x, p.y, p.z] for p in pores])
    if len(coords) > 0:
        print(f"Bounds: Min {coords.min(axis=0)}, Max {coords.max(axis=0)}")
        
    return pores
 
def save_vtp(filename, pores):
    import struct
    with open(filename, 'w') as f:
        f.write('<?xml version="1.0"?>\n')
        f.write('<VTKFile type="PolyData" version="0.1" byte_order="LittleEndian">\n')
        f.write('  <PolyData>\n')
        f.write(f'    <Piece NumberOfPoints="{len(pores)}" NumberOfVerts="0" NumberOfLines="0" NumberOfStrips="0" NumberOfPolys="0">\n')
        
        f.write('      <Points>\n')
        f.write('        <DataArray type="Float32" NumberOfComponents="3" format="ascii">\n')
        for p in pores:
            f.write(f'{p.x} {p.y} {p.z}\n')
        f.write('        </DataArray>\n')
        f.write('      </Points>\n')
        
        f.write('      <PointData Scalars="Radius">\n')
        f.write('        <DataArray type="Float32" Name="Radius" format="ascii">\n')
        for p in pores:
            f.write(f'{p.radius}\n')
        f.write('        </DataArray>\n')
        f.write('      </PointData>\n')
        
        f.write('    </Piece>\n')
        f.write('  </PolyData>\n')
        f.write('</VTKFile>\n')
    print(f"Saved pores to {filename}")
 
if __name__ == "__main__":
    if len(sys.argv) > 1:
        pores = test_extraction(sys.argv[1])
        save_vtp("pores.vtp", pores)
    else:
        print(f"Usage: {sys.argv[0]} <vti_file>")
        sys.exit(2)