Special case for writing 3d OFF

3 files changed, 28 insertions, 1 deletions

diff --git a/src/python/CMakeLists.txt b/src/python/CMakeLists.txt
index 35ddb778..32ec13bd 100644
--- a/src/python/CMakeLists.txt
+++ b/src/python/CMakeLists.txt
@@ -546,6 +546,7 @@ if(PYTHONINTERP_FOUND)
 
     # Reader utils
     add_gudhi_py_test(test_reader_utils)
+    add_gudhi_py_test(test_off)
 
     # Wasserstein
     if(OT_FOUND)
diff --git a/src/python/gudhi/off_utils.pyx b/src/python/gudhi/off_utils.pyx
index a8142791..9276c7b0 100644
--- a/src/python/gudhi/off_utils.pyx
+++ b/src/python/gudhi/off_utils.pyx
@@ -54,4 +54,9 @@ def write_points_to_off_file(fname, points):
     """
     points = np.array(points, copy=False)
     assert len(points.shape) == 2
-    np.savetxt(fname, points, header='nOFF\n{} {} 0 0'.format(points.shape[1], points.shape[0]), comments='')
+    dim = points.shape[1]
+    if dim == 3:
+        head = 'OFF\n{} 0 0'.format(points.shape[0])
+    else:
+        head = 'nOFF\n{} {} 0 0'.format(dim, points.shape[0])
+    np.savetxt(fname, points, header=head, comments='')
diff --git a/src/python/test/test_off.py b/src/python/test/test_off.py
new file mode 100644
index 00000000..69bfa1f9
--- /dev/null
+++ b/src/python/test/test_off.py
@@ -0,0 +1,21 @@
+""" This file is part of the Gudhi Library - https://gudhi.inria.fr/ - which is released under MIT.
+    See file LICENSE or go to https://gudhi.inria.fr/licensing/ for full license details.
+    Author(s):       Marc Glisse
+
+    Copyright (C) 2022 Inria
+
+    Modification(s):
+      - YYYY/MM Author: Description of the modification
+"""
+
+import gudhi as gd
+import numpy as np
+import pytest
+
+
+def test_off_rw():
+    for dim in range(2, 6):
+        X = np.random.rand(123, dim)
+        gd.write_points_to_off_file('rand.off', X)
+        Y = gd.read_points_from_off_file('rand.off')
+        assert Y == pytest.approx(X)