From 2d9eb1826ccc9b4276448ba65c7550e1d7cb3843 Mon Sep 17 00:00:00 2001
From: Axel Larsson <65452706+axla-io@users.noreply.github.com>
Date: Wed, 30 Oct 2024 14:39:44 -0400
Subject: [PATCH] Improved graph handling

---
 src/DynamicRelaxation.jl          |  4 ++--
 src/analysis/3dof_acceleration.jl |  3 ++-
 src/analysis/graph.jl             | 34 +++++++++++++++++++------------
 src/optimization/load_finding.jl  |  4 ++--
 src/simulation.jl                 |  5 +++--
 5 files changed, 30 insertions(+), 20 deletions(-)

diff --git a/src/DynamicRelaxation.jl b/src/DynamicRelaxation.jl
index e65e742..c9311da 100644
--- a/src/DynamicRelaxation.jl
+++ b/src/DynamicRelaxation.jl
@@ -2,7 +2,7 @@ module DynamicRelaxation
 
 # Arrays
 using LinearAlgebra
-using SparseArrays: sparse, spzeros, findnz
+using SparseArrays: sparse, spzeros, findnz, nzrange, rowvals
 using StaticArrays
 using ConcreteStructs
 using UnPack
@@ -47,7 +47,7 @@ export BigonTorqueCondition, clamped, free, pinned, roller
 export Px, Py, Pz, Mx, My, Mz, uniform_load, point_loads
 
 # System 
-export StructuralGraphSystem, default_system, get_cs, get_ode_jac, create_graph
+export StructuralGraphSystem, default_system, get_cs, get_ode_jac, create_graph, edges
 
 # Simulation
 export LoadScaleRodSimulation, RodSimulation, BigonRodSimulation, get_u0, get_vel_ids, get_state
diff --git a/src/analysis/3dof_acceleration.jl b/src/analysis/3dof_acceleration.jl
index 3d0e227..73584e9 100644
--- a/src/analysis/3dof_acceleration.jl
+++ b/src/analysis/3dof_acceleration.jl
@@ -6,7 +6,8 @@ function rod_acceleration(x, system::StructuralGraphSystem{Vector{Node3DOF}}, vt
 	a = @SVector zeros(eltype(x), 3)
 	s = @SVector zeros(eltype(x), 3)
 
-	for neighbor in neighbors(graph, vtx)
+	for idx in nzrange(graph, vtx)
+		neighbor = rowvals(graph)[idx]
 		n_i = 3 * (neighbor - 1) + 1
 		ep = elem_props[edge_index((vtx, neighbor), graph)]
 		(a, s) = rod_accelerate(a, x_vert, @view(x[n_i:(n_i+2)]), ep, s)
diff --git a/src/analysis/graph.jl b/src/analysis/graph.jl
index 4e313af..f63783f 100644
--- a/src/analysis/graph.jl
+++ b/src/analysis/graph.jl
@@ -1,28 +1,36 @@
 function neighbors(graph, vtx)
-	neighbors = findnz(graph[vtx, :])[2]
+	neighbors = findnz(graph[vtx, :])[1]
 	return neighbors
 end
 
 function edges(graph)
-	return findnz(graph)
+    is, js, ids = findnz(graph)
+
+    # Filter list
+    pos_ids = ids .>= 0.0
+    is_pos = is[pos_ids]
+    js_pos = js[pos_ids]
+    ids_pos = ids[pos_ids]
+
+    # Order list
+    order_dict = Dict(v => i for (i, v) in enumerate(ids_pos))
+
+    ids_sort = sortperm(ids_pos, by=x -> order_dict[x])
+    is_sort = is_pos[ids_sort]
+    js_sort = js_pos[ids_sort]
+	return zip(is_sort, js_sort)
 end
 
 function create_graph(edges, pts)
     n_pts = length(pts)
 	graph = spzeros(Int, n_pts, n_pts)
-	for edge in edges
-		graph[edge[1], edge[2]] = 1
-		graph[edge[2], edge[1]] = 1
+	for (i, edge) in enumerate(edges)
+		graph[edge[1], edge[2]] = i
+		graph[edge[2], edge[1]] = -i
 	end
 	return graph
 end
 
-function edge_index(adj_matrix, i, j)
-	row_indices, col_indices, _ = findnz(adj_matrix)
-	for idx in 1:length(row_indices)
-		if row_indices[idx] == i && col_indices[idx] == j
-			return idx  # Position of the edge in the sparse storage
-		end
-	end
-	return nothing  # Edge not found
+function edge_index(edge, graph)
+	return abs(graph[edge[1], edge[2]])
 end
diff --git a/src/optimization/load_finding.jl b/src/optimization/load_finding.jl
index 6c6fe1b..495e30b 100644
--- a/src/optimization/load_finding.jl
+++ b/src/optimization/load_finding.jl
@@ -21,7 +21,7 @@ function f_acceleration(a, ext_f, i::Int, p)
 end
 
 function accelerate_system(u_v, system::StructuralGraphSystem{Vector{Node6DOF}},
-                           simulation::LoadScaleRodSimulation{Node6DOF}, body,
+                           simulation::LoadScaleRodSimulation{Vector{Node6DOF}}, body,
                            ext_f, du, dr_ids, ω, i, dt, u_t, p, t)
     (a, τ, s, j) = rod_acceleration(u_v, system, body, i)
     (a, τ) = f_acceleration(a, τ, ext_f, i, p)
@@ -32,7 +32,7 @@ function accelerate_system(u_v, system::StructuralGraphSystem{Vector{Node6DOF}},
 end
 
 function accelerate_system(u_v, system::StructuralGraphSystem{Vector{Node3DOF}},
-                           simulation::LoadScaleRodSimulation{Node3DOF}, body,
+                           simulation::LoadScaleRodSimulation{Vector{Node3DOF}}, body,
                            ext_f, du, dr_ids, ω, i, dt, u_t, p, t)
     (a, s) = rod_acceleration(u_v, system, i)
     rod_acceleration(u_v, system, i)
diff --git a/src/simulation.jl b/src/simulation.jl
index c51ffad..1927435 100644
--- a/src/simulation.jl
+++ b/src/simulation.jl
@@ -230,8 +230,9 @@ function accelerate_system(u_v, system::StructuralGraphSystem{Vector{Node6DOF}},
 end
 
 function accelerate_system(u_v, system::StructuralGraphSystem{Vector{Node3DOF}},
-	simulation::RodSimulation{Node3DOF}, body,
+	simulation::RodSimulation{Vector{Node3DOF}}, body,
 	ext_f, du, dr_ids, ω, i, dt, u_t, p, t)
+	rod_acceleration(u_v, system, i)
 	(a, s) = rod_acceleration(u_v, system, i)
 	a = f_acceleration(a, ext_f, i)
 	a = constrain_acceleration(a, body)
@@ -240,7 +241,7 @@ function accelerate_system(u_v, system::StructuralGraphSystem{Vector{Node3DOF}},
 end
 
 function get_system_forces(u_v, system::StructuralGraphSystem{Vector{Node6DOF}},
-	simulation::RodSimulation{Node6DOF}, body,
+	simulation::RodSimulation{Vector{Node6DOF}}, body,
 	ext_f, du, dr_ids, ω, i, dt, u_t, p)
 	(a, τ, s, j) = rod_acceleration(u_v, system, body, i)