2d sort

dsmc/common.py

@@ -1,5 +1,11 @@
-from numba import cfunc
+from numba import cfunc, njit
 
 @cfunc("double(double[::3, ::2])")
 def get_V(a):
-    return (a[0][1] - a[0][0]) * (a[1][1] - a[1][0]) * (a[2][1] - a[2][0])
+    return (a[0][1] - a[0][0]) * (a[1][1] - a[1][0]) * (a[2][1] - a[2][0])
+
+@njit 
+def swap(arr, pos1, pos2):
+    temp = arr[pos2]
+    arr[pos2] = arr[pos1]
+    arr[pos1] = temp

dsmc/diagnostics.py

@@ -1,6 +1,67 @@
 import numpy as np
+from numba import njit
 from . import particles as prt
 from . import octree as oc
+from . import common as com
+
+@njit
+def _in_2d_box(pos, box):
+    a = pos[0] >= box[0][0] and pos[0] <= box[0][1]
+    b = pos[1] >= box[1][0] and pos[1] <= box[1][1]
+    return a and b
+
+@njit
+def _set_up_2d_boxes(x0, x1, y0, y1, Nx, Ny):
+    dx = (x1 - x0) / Nx
+    dy = (y1 - y0) / Ny
+    boxes = np.empty((Nx*Ny, 2, 2))
+    y = y0
+
+    for i in range(Ny):
+        x = x0
+        for j in range(Nx):
+            k = j + i*Nx
+            boxes[k][0][0] = x
+            boxes[k][0][1] = x + dx
+            boxes[k][1][0] = y
+            boxes[k][1][1] = y + dy
+            x += dx
+        y += dy
+
+    return boxes
+
+@njit
+def _sort_2d(positions, x0, x1, y0, y1, Nx, Ny):
+    boxes = _set_up_2d_boxes(x0, x1, y0, y1, Nx, Ny) # [] : bix id, [][] : axis, [][][] : 0 min , 1 max
+    permutations = [i for i in range(len(positions))]
+    leafs = np.zeros((Nx * Ny, 2), dtype=np.int_) # leaf[0] : ofseat, leaf[1] : n-parts
+
+    for i in range(len(leafs)):
+        leafs[i][0] = leafs[i - 1][0] + leafs[i - 1][1] if i > 0 else 0
+        leafs[i][1] = 0
+        runner = leafs[i][0]
+        for j in range(leafs[i][0], len(positions)):
+            p = permutations[j]
+            if _in_2d_box(positions[p], boxes[i]):
+                com.swap(permutations, j, runner)
+                runner += 1
+                leafs[i][1] += 1
+
+    return (permutations, leafs, boxes)
+
+class Values:
+    def __init__(self):
+        self.number_elements = 0
+
+def analyse_2d(positions, x0, x1, y0, y1, Nx, Ny):
+    permutations, leafs, boxes = _sort_2d(positions, x0, x1, y0, y1, Nx, Ny)
+    values = [Values() for _ in range(len(leafs))]
+
+    for i in range(len(leafs)):
+        values[i].number_elements = leafs[i][1]
+
+    return (boxes, values)
+
 
 def sort_bin(positions, axis, Nbin):
     bins = [[] for _ in range(Nbin)]

dsmc/octree.py

@@ -1,6 +1,7 @@
 import numpy as np
 import numpy.typing as npt
 from numba import njit
+from . import common as com
 
 fmin = np.finfo(float).min
 fmax = np.finfo(float).max
@@ -69,12 +70,6 @@ def _is_inside(position : npt.NDArray, box : npt.NDArray) -> bool:
     c : bool = position[2] >= box[2][0] and position[2] <= box[2][1]
 
     return a and b and c
-
-@njit    
-def _swap(arr, pos1, pos2):
-    temp = arr[pos2]
-    arr[pos2] = arr[pos1]
-    arr[pos1] = temp
 
 @njit
 def _sort(permutations : npt.NDArray, box : npt.NDArray, positions : npt.NDArray, offset : int, N : int) -> tuple[npt.NDArray, int]:
@@ -103,7 +98,7 @@ def _sort(permutations : npt.NDArray, box : npt.NDArray, positions : npt.NDArray
     for i in range(offset, offset + N):
         p = new_permutations[i]
         if _is_inside(positions[p], box):
-            _swap(new_permutations, i, runner)
+            com.swap(new_permutations, i, runner)
             runner += 1
             Nnew += 1
 

dsmc/writer/__init__.py

@@ -0,0 +1,2 @@
+from .octree import write_buttom_leafs
+from .planar import write as write_planar

dsmc/writer/planar.py

@@ -0,0 +1,86 @@
+def write(boxes, values, file_name="planar.vtu"):
+	f = open(file_name, "w")
+
+	_write_header(f)
+	_wrtie_body(f, boxes, values)
+	_write_footer(f)
+
+	f.close()
+
+def _write_header(f):
+	f.write("<VTKFile type=\"UnstructuredGrid\" version=\"0.1\" byte_order=\"LittleEndian\">\n")
+
+def _wrtie_body(f, boxes, values):
+	f.write("<UnstructuredGrid>\n")
+	f.write("<Piece NumberOfPoints=\"{}\" NumberOfCells=\"{}\">\n".format(len(boxes) * 4, len(boxes)))
+
+	_write_points(f, boxes)
+	_write_cells(f, boxes)
+	_write_cell_data(f, boxes, values)
+
+	f.write("</Piece>\n")
+	f.write("</UnstructuredGrid>\n")
+
+def _write_points(f, boxes):
+	f.write("<Points>\n")
+	f.write("<DataArray type=\"Float32\" NumberOfComponents=\"3\" format=\"ascii\">\n")
+
+	for box in boxes:
+		f.write("{} ".format(box[0][0]))
+		f.write("{} ".format(box[1][0]))
+		f.write("{} ".format(0.0))
+
+		f.write("{} ".format(box[0][1]))
+		f.write("{} ".format(box[1][0]))
+		f.write("{} ".format(0.0))
+
+		f.write("{} ".format(box[0][1]))
+		f.write("{} ".format(box[1][1]))
+		f.write("{} ".format(0.0))
+
+		f.write("{} ".format(box[0][0]))
+		f.write("{} ".format(box[1][1]))
+		f.write("{} ".format(0.0))
+
+	f.write("</DataArray>\n")
+	f.write("</Points>\n")
+
+
+def _write_cells(f, boxes):
+	k = 0
+
+	f.write("<Cells>\n")
+	f.write("<DataArray type=\"Int32\" Name=\"connectivity\" format=\"ascii\">\n")
+
+	for i in range(len(boxes)):
+		for _ in range(4):
+			f.write("{} ".format(k))
+			k += 1
+
+	f.write("</DataArray>\n")
+	f.write("<DataArray type=\"Int32\" Name=\"offsets\" format=\"ascii\">\n")
+
+	for i in range(len(boxes)):
+		f.write("{} ".format((i + 1) * 4))
+
+	f.write("</DataArray>\n")
+	f.write("<DataArray type=\"UInt8\" Name=\"types\" format=\"ascii\">\n")
+
+	for _ in range(len(boxes)):
+		f.write("9 ")
+
+	f.write("</DataArray>\n")
+	f.write("</Cells>\n")
+
+def _write_cell_data(f, boxes, values):
+    f.write("<CellData Scalars=\"particle_numbers\">\n")
+    f.write("<DataArray type=\"Float32\" Name=\"particle_numbers\" format=\"ascii\">\n")
+
+    for vals in values:
+        f.write("{} ".format(vals.number_elements))
+
+    f.write("</DataArray>\n")
+    f.write("</CellData>\n")
+
+def _write_footer(f):
+	f.write("</VTKFile>\n")

setup.py

@@ -4,7 +4,7 @@
     version='0.7.0',
     author='Leo Basov',
     python_requires='>=3.10, <4',
-    packages=["dsmc"],
+    packages=["dsmc", "dsmc.writer"],
     install_requires=[
         'numpy',
         'llvmlite',

tests/diagnostics/dia_test.py

@@ -0,0 +1,45 @@
+import numpy as np
+import dsmc.diagnostics as dia
+import dsmc.writer as wrt
+import time
+
+def create_particles(N, radius):
+	positions = np.empty((N + 1, 3))
+
+	for i in range(N):
+		phi = np.random.random() * 2.0 * np.pi
+		theta = np.random.random() * np.pi
+		r = np.random.normal(0.0, 0.01)
+		theta1 = np.random.random() * np.pi - 0.5 * np.pi
+		dis = np.array((r * np.sin(theta) * np.cos(phi), r * np.sin(theta) * np.sin(phi), r * np.cos(theta)))
+		offset = np.array((radius * np.sin(theta1), radius * np.cos(theta1), 0.0))
+
+		dis += offset;
+		positions[i] = dis
+
+	positions[N] = np.array((0.0, -1.0, 0.0))
+
+	return positions
+
+if __name__ == "__main__":
+    N = 10000
+    radius = 1.0
+    positions = create_particles(N, radius)
+
+    # set up sort
+    Nx = 100
+    Ny = 100
+    x0 = -1.5
+    x1 = 1.5
+    y0 = -1.5
+    y1 = 1.5
+
+    start_time = time.time()
+    boxes, values = dia.analyse_2d(positions, x0, x1, y0, y1, Nx, Ny)
+    print("--- %s seconds ---" % (time.time() - start_time))   
+
+    print("writing to file")
+
+    wrt.write_planar(boxes, values)
+
+    print("done")