Merge pull request #403 from jcapriot/TreeMesh_modify_lock

Updates for python free threading

Author: jcapriot

.ci/azure/setup_env.sh

@@ -4,6 +4,7 @@ set -ex #echo on and exit if any line fails
 # TF_BUILD is set to True on azure pipelines.
 is_azure=$(echo "${TF_BUILD:-false}" | tr '[:upper:]' '[:lower:]')
 do_doc=$(echo "${DOC_BUILD:-false}" | tr '[:upper:]' '[:lower:]')
+is_free_threaded=$(echo "${PYTHON_FREETHREADING:-false}" | tr '[:upper:]' '[:lower:]')
 
 if ${is_azure}
 then
@@ -13,8 +14,14 @@ then
   fi
 fi
 
-cp .ci/environment_test.yml environment_test_with_pyversion.yml
-echo "  - python="$PYTHON_VERSION >> environment_test_with_pyversion.yml
+if ${is_free_threaded}
+then
+  cp .ci/environment_test_bare.yml environment_test_with_pyversion.yml
+  echo "  - python-freethreading="$PYTHON_VERSION >> environment_test_with_pyversion.yml
+else
+  cp .ci/environment_test.yml environment_test_with_pyversion.yml
+  echo "  - python="$PYTHON_VERSION >> environment_test_with_pyversion.yml
+fi
 
 conda env create --file environment_test_with_pyversion.yml
 rm environment_test_with_pyversion.yml

.ci/azure/test.yml

@@ -5,16 +5,21 @@ jobs:
       linux-Python310:
         image: ubuntu-latest
         python.version: '3.10'
-        coverage: True
       linux-Python311:
         image: ubuntu-latest
         python.version: '3.11'
+        coverage: True
       linux-Python312:
         image: ubuntu-latest
         python.version: '3.12'
       linux-Python313:
         image: ubuntu-latest
         python.version: '3.13'
+      linux-Python313t:
+        image: ubuntu-latest
+        python.version: '3.13'
+        python.freethreading: True
+        coverage: True
       osx-Python310:
         image: macOS-latest
         python.version: '3.10'
@@ -27,6 +32,10 @@ jobs:
       osx-Python313:
         image: macOS-latest
         python.version: '3.13'
+      osx-Python313t:
+        image: macOS-latest
+        python.version: '3.13'
+        python.freethreading: True
       win-Python310:
         image: windows-latest
         python.version: '3.10'
@@ -39,6 +48,10 @@ jobs:
       win-Python313:
         image: windows-latest
         python.version: '3.13'
+      win-Python313t:
+        image: windows-latest
+        python.version: '3.13'
+        python.freethreading: True
   displayName: "${{ variables.image }} ${{ variables.python.version }}"
   pool:
     vmImage: $(image)

.ci/environment_test.yml

@@ -31,5 +31,5 @@ dependencies:
   - meson-python>=0.14.0
   - meson
   - ninja
-  - cython>=0.29.35
+  - cython>=3.1.0
   - setuptools_scm

.ci/environment_test_bare.yml

@@ -0,0 +1,19 @@
+name: discretize-test
+channels:
+  - conda-forge
+dependencies:
+  - numpy>=1.22.4
+  - scipy>=1.8
+
+  # testing
+  - sympy
+  - pytest
+  - pytest-cov
+
+  # Building
+  - pip
+  - meson-python>=0.14.0
+  - meson
+  - ninja
+  - cython>=3.1.0
+  - setuptools_scm

discretize/_extensions/interputils_cython.pyx

@@ -1,5 +1,6 @@
 # cython: embedsignature=True, language_level=3
 # cython: linetrace=True
+# cython: freethreading_compatible = True
 import numpy as np
 import cython
 cimport numpy as np

discretize/_extensions/simplex_helpers.pyx

@@ -1,6 +1,7 @@
 # distutils: language=c++
 # cython: embedsignature=True, language_level=3
 # cython: linetrace=True
+# cython: freethreading_compatible = True
 
 from libcpp.pair cimport pair
 from libcpp.unordered_map cimport unordered_map

discretize/_extensions/tree_ext.pyx

@@ -1,6 +1,7 @@
 # distutils: language=c++
 # cython: embedsignature=True, language_level=3
 # cython: linetrace=True
+# cython: freethreading_compatible=True
 cimport cython
 cimport numpy as np
 from libc.stdlib cimport malloc, free
@@ -359,6 +360,7 @@ cdef class _TreeMesh:
     cdef int_t[3] ls
     cdef int _finalized
     cdef bool _diagonal_balance
+    cdef cython.pymutex _tree_modify_lock
 
     cdef double[:] _xs, _ys, _zs
     cdef double[:] _origin
@@ -546,9 +548,12 @@ cdef class _TreeMesh:
 
         #Wrapping function so it can be called in c++
         cdef void * func_ptr = <void *> function
-        self.wrapper.set(func_ptr, _evaluate_func)
-        #Then tell c++ to build the tree
-        self.tree.refine_function(self.wrapper, diag_balance)
+        
+
+        with self._tree_modify_lock:
+            self.wrapper.set(func_ptr, _evaluate_func)
+            #Then tell c++ to build the tree
+            self.tree.refine_function(self.wrapper, diag_balance)
         if finalize:
             self.finalize()
 
@@ -631,7 +636,9 @@ cdef class _TreeMesh:
         for i in range(n_balls):
             ball = geom.Ball(self._dim, &cs[i_c, 0], rs[i_r])
             l = _wrap_levels(ls[i_l], max_level)
-            self.tree.refine_geom(ball, l, diag_balance)
+        
+            with self._tree_modify_lock:
+                    self.tree.refine_geom(ball, l, diag_balance)
 
             i_c += cs_step
             i_r += rs_step
@@ -715,7 +722,8 @@ cdef class _TreeMesh:
         for i in range(n_boxes):
             box = geom.Box(self._dim, &x0[i_x0, 0], &x1[i_x1, 0])
             l = _wrap_levels(ls[i_l], max_level)
-            self.tree.refine_geom(box, l, diag_balance)
+            with self._tree_modify_lock:
+                self.tree.refine_geom(box, l, diag_balance)
 
             i_x0 += x0_step
             i_x1 += x1_step
@@ -792,7 +800,8 @@ cdef class _TreeMesh:
         for i in range(n_segments):
             line = geom.Line(self._dim, &line_nodes[i_line, 0], &line_nodes[i_line+1, 0])
             l = _wrap_levels(ls[i_l], max_level)
-            self.tree.refine_geom(line, l, diag_balance)
+            with self._tree_modify_lock:
+                self.tree.refine_geom(line, l, diag_balance)
 
             i_line += line_step
             i_l += l_step
@@ -875,7 +884,8 @@ cdef class _TreeMesh:
         for i in range(n_planes):
             plane = geom.Plane(self._dim, &x_0s[i_o, 0], &norms[i_n, 0])
             l = _wrap_levels(ls[i_l], max_level)
-            self.tree.refine_geom(plane, l, diag_balance)
+            with self._tree_modify_lock:
+                self.tree.refine_geom(plane, l, diag_balance)
 
             i_o += origin_step
             i_n += normal_step
@@ -953,7 +963,8 @@ cdef class _TreeMesh:
         for i in range(n_triangles):
             triang = geom.Triangle(self._dim, &tris[i_tri, 0, 0], &tris[i_tri, 1, 0], &tris[i_tri, 2, 0])
             l = _wrap_levels(ls[i_l], max_level)
-            self.tree.refine_geom(triang, l, diag_balance)
+            with self._tree_modify_lock:
+                self.tree.refine_geom(triang, l, diag_balance)
 
             i_tri += tri_step
             i_l += l_step
@@ -1048,7 +1059,8 @@ cdef class _TreeMesh:
         for i in range(n_triangles):
             vert_prism = geom.VerticalTriangularPrism(self._dim, &tris[i_tri, 0, 0], &tris[i_tri, 1, 0], &tris[i_tri, 2, 0], hs[i_h])
             l = _wrap_levels(ls[i_l], max_level)
-            self.tree.refine_geom(vert_prism, l, diag_balance)
+            with self._tree_modify_lock:
+                self.tree.refine_geom(vert_prism, l, diag_balance)
 
             i_tri += tri_step
             i_h += h_step
@@ -1133,7 +1145,9 @@ cdef class _TreeMesh:
         for i in range(n_triangles):
             l = _wrap_levels(ls[i_l], max_level)
             tet = geom.Tetrahedron(self._dim, &tris[i_tri, 0, 0], &tris[i_tri, 1, 0], &tris[i_tri, 2, 0], &tris[i_tri, 3, 0])
-            self.tree.refine_geom(tet, l, diag_balance)
+            
+            with self._tree_modify_lock:
+                self.tree.refine_geom(tet, l, diag_balance)
 
             i_tri += tri_step
             i_l += l_step
@@ -1192,7 +1206,8 @@ cdef class _TreeMesh:
 
         for i in range(ls.shape[0]):
             l = _wrap_levels(ls[i_l], max_level)
-            self.tree.insert_cell(&cs[i_p, 0], l, diagonal_balance)
+            with self._tree_modify_lock:
+                self.tree.insert_cell(&cs[i_p, 0], l, diagonal_balance)
 
             i_l += l_step
             i_p += p_step
@@ -1208,10 +1223,11 @@ cdef class _TreeMesh:
         operators. When finalized, mesh refinement is no longer enabled.
 
         """
-        if not self._finalized:
-            self.tree.finalize_lists()
-            self.tree.number()
-            self._finalized=True
+        with self._tree_modify_lock:
+            if not self._finalized:
+                self.tree.finalize_lists()
+                self.tree.number()
+                self._finalized=True
 
     @property
     def finalized(self):
@@ -1228,7 +1244,9 @@ cdef class _TreeMesh:
         bool
             Returns *True* if finalized, *False* otherwise
         """
-        return self._finalized
+        with self._tree_modify_lock:
+            val = self._finalized
+        return val
 
     @property
     @cython.boundscheck(False)
@@ -1248,7 +1266,9 @@ cdef class _TreeMesh:
 
     def number(self):
         """Number the cells, nodes, faces, and edges of the TreeMesh."""
-        self.tree.number()
+        
+        with self._tree_modify_lock:
+            self.tree.number()
 
     def get_containing_cells(self, points):
         """Return the cells containing the given points.
@@ -1453,52 +1473,53 @@ cdef class _TreeMesh:
             if dim == 3:
                 shift[2] = self._origin[2] - self._zs[0]
 
-            for i in range(self._xs.shape[0]):
-                self._xs[i] += shift[0]
-            for i in range(self._ys.shape[0]):
-                self._ys[i] += shift[1]
-            if dim == 3:
-                for i in range(self._zs.shape[0]):
-                    self._zs[i] += shift[2]
-
-            #update the locations of all of the items
-            self.tree.shift_cell_centers(&shift[0])
-
-            for itN in self.tree.nodes:
-                node = itN.second
-                for i in range(dim):
-                    node.location[i] += shift[i]
+            with self._tree_modify_lock:
+                for i in range(self._xs.shape[0]):
+                    self._xs[i] += shift[0]
+                for i in range(self._ys.shape[0]):
+                    self._ys[i] += shift[1]
+                if dim == 3:
+                    for i in range(self._zs.shape[0]):
+                        self._zs[i] += shift[2]
 
-            for itE in self.tree.edges_x:
-                edge = itE.second
-                for i in range(dim):
-                    edge.location[i] += shift[i]
+                #update the locations of all of the items
+                self.tree.shift_cell_centers(&shift[0])
 
-            for itE in self.tree.edges_y:
-                edge = itE.second
-                for i in range(dim):
-                    edge.location[i] += shift[i]
+                for itN in self.tree.nodes:
+                    node = itN.second
+                    for i in range(dim):
+                        node.location[i] += shift[i]
 
-            if dim == 3:
-                for itE in self.tree.edges_z:
+                for itE in self.tree.edges_x:
                     edge = itE.second
                     for i in range(dim):
                         edge.location[i] += shift[i]
 
-                for itF in self.tree.faces_x:
-                    face = itF.second
-                    for i in range(dim):
-                        face.location[i] += shift[i]
-
-                for itF in self.tree.faces_y:
-                    face = itF.second
+                for itE in self.tree.edges_y:
+                    edge = itE.second
                     for i in range(dim):
-                        face.location[i] += shift[i]
+                        edge.location[i] += shift[i]
 
-                for itF in self.tree.faces_z:
-                    face = itF.second
-                    for i in range(dim):
-                        face.location[i] += shift[i]
+                if dim == 3:
+                    for itE in self.tree.edges_z:
+                        edge = itE.second
+                        for i in range(dim):
+                            edge.location[i] += shift[i]
+
+                    for itF in self.tree.faces_x:
+                        face = itF.second
+                        for i in range(dim):
+                            face.location[i] += shift[i]
+
+                    for itF in self.tree.faces_y:
+                        face = itF.second
+                        for i in range(dim):
+                            face.location[i] += shift[i]
+
+                    for itF in self.tree.faces_z:
+                        face = itF.second
+                        for i in range(dim):
+                            face.location[i] += shift[i]
             #clear out all cached grids
             self._cell_centers = None
             self._nodes = None
@@ -7046,7 +7067,6 @@ cdef class _TreeMesh:
             )
         return np.require(arr, dtype=dtype, requirements=requirements)
 
-
 def _check_first_dim_broadcast(**kwargs):
     """Perform a check to make sure that the first dimensions of the inputs will broadcast."""
     n_items = 1

pyproject.toml

@@ -3,7 +3,7 @@
 build-backend = 'mesonpy'
 requires = [
     "meson-python>=0.15.0",
-    "Cython>=3.0.8",
+    "Cython>=3.1.0",
     "setuptools_scm[toml]>=6.2",
 
     # numpy requirement for wheel builds for distribution on PyPI - building
@@ -93,7 +93,7 @@ build = [
     "meson",
     "ninja",
     "numpy>=1.22.4",
-    "cython>=0.29.35",
+    "cython>=3.1.0",
     "setuptools_scm",
 ]
 
@@ -124,9 +124,9 @@ setup-args = [
 [tool.coverage.run]
 branch = true
 source = ["discretize", "tests", "examples", "tutorials"]
-plugins = [
-    "Cython.Coverage",
-]
+# plugins = [
+#     "Cython.Coverage",
+# ]
 
 [tool.coverage.report]
 ignore_errors = false