From 9c029bc2f53abd582aa1f44f31717eb4c0d12573 Mon Sep 17 00:00:00 2001
From: tomvanmele <tom.v.mele@gmail.com>
Date: Wed, 28 Feb 2024 22:38:54 +0100
Subject: [PATCH] test scripts

---
 scripts/test_dr_constrained_numpy.py | 83 ++++++++++++++++++++++++++++
 scripts/test_dr_numpy.py             |  2 +-
 2 files changed, 84 insertions(+), 1 deletion(-)
 create mode 100644 scripts/test_dr_constrained_numpy.py

diff --git a/scripts/test_dr_constrained_numpy.py b/scripts/test_dr_constrained_numpy.py
new file mode 100644
index 0000000..68484e5
--- /dev/null
+++ b/scripts/test_dr_constrained_numpy.py
@@ -0,0 +1,83 @@
+from compas.colors import Color
+from compas.datastructures import Mesh
+from compas.geometry import Line
+from compas.geometry import NurbsCurve
+from compas.geometry import Point
+from compas.geometry import Sphere
+from compas.geometry import Vector
+from compas_dr.constraints import Constraint
+from compas_dr.numdata import InputData
+from compas_dr.solvers import dr_constrained_numpy
+from compas_view2.app import App
+
+# =============================================================================
+# Input
+# =============================================================================
+
+mesh = Mesh.from_meshgrid(dx=10, nx=10)
+
+fixed = list(mesh.vertices_where(vertex_degree=2))
+loads = [[0, 0, 0]] * mesh.number_of_vertices()
+qpre = [1.0] * mesh.number_of_edges()
+
+for index, edge in enumerate(mesh.edges()):
+    if mesh.is_edge_on_boundary(edge):
+        qpre[index] = 10
+
+qpre = []
+for edge in mesh.edges():
+    if mesh.is_edge_on_boundary(edge):
+        qpre.append(10)
+    else:
+        qpre.append(1.0)
+
+inputdata = InputData.from_mesh(mesh, fixed, loads, qpre)
+
+# constraints
+
+arch = NurbsCurve.from_points([[5, 0, 0], [5, 5, 5], [5, 10, 0]])
+constraint = Constraint(arch)
+
+constraints = [None] * mesh.number_of_vertices()
+for vertex in mesh.vertices_where(x=5):
+    constraints[vertex] = constraint
+    fixed.append(vertex)
+
+# =============================================================================
+# Solve and Update
+# =============================================================================
+
+result = dr_constrained_numpy(indata=inputdata, constraints=constraints, kmax=1000)
+
+for vertex in mesh.vertices():
+    mesh.vertex_attributes(vertex, "xyz", result.xyz[vertex])
+
+# =============================================================================
+# Visualization
+# =============================================================================
+
+viewer = App()
+viewer.add(mesh)
+
+for vertex in mesh.vertices():
+    point = Point(*mesh.vertex_coordinates(vertex))
+    residual = Vector(*result.residuals[vertex])
+
+    if vertex in fixed:
+        ball = Sphere(radius=0.1, point=point)
+
+        if constraints[vertex]:
+            viewer.add(ball.to_brep(), facecolor=Color.blue())
+        else:
+            viewer.add(ball.to_brep(), facecolor=Color.red())
+
+        viewer.add(
+            Line(point, point - residual * 0.1),
+            linecolor=Color.green().darkened(50),
+            linewidth=3,
+        )
+
+viewer.add(arch.to_polyline(), linecolor=Color.cyan(), linewidth=3)
+
+viewer.view.camera.zoom_extents()
+viewer.run()
diff --git a/scripts/test_dr_numpy.py b/scripts/test_dr_numpy.py
index 23b46a6..4d45507 100644
--- a/scripts/test_dr_numpy.py
+++ b/scripts/test_dr_numpy.py
@@ -4,8 +4,8 @@
 from compas.geometry import Point
 from compas.geometry import Sphere
 from compas.geometry import Vector
-from compas_dr.dr_numpy import dr_numpy
 from compas_dr.numdata import InputData
+from compas_dr.solvers.dr_numpy import dr_numpy
 from compas_view2.app import App
 
 # =============================================================================