Unit testing

Author: Chad-Peterson

examples/Scratch/global_test.py

@@ -28,24 +28,6 @@
       | /            | /
       4--------------5
 
-* | 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7
-----------------------------------
-0 |   |   |   |   |   |   |   |
-----------------------------------
-1 |   |   |   |   |   |   |   |
-----------------------------------
-2 |   |   |   |   |   |   |   |
-----------------------------------
-3 |   |   |   |   |   |   |   |
-----------------------------------
-4 |   |   |   |   |   |   |   |
-----------------------------------
-5 |   |   |   |   |   |   |   |
-----------------------------------
-6 |   |   |   |   |   |   |   |
-----------------------------------
-7 |   |   |   |   |   |   |   |
-
 A two-element mesh:
 
              Element 1       Element 2
@@ -59,33 +41,6 @@
     | /            | /            | /
     6--------------7--------------8
 
-
-* | 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11
----------------------------------------------------
-0 |e1 |e1 |   |   |   |   |   |   |   |   |    |
----------------------------------------------------
-1 |e1 |e12|   |   |   |   |   |   |   |   |    |
----------------------------------------------------
-2 |   |   |e2 |   |   |   |   |   |   |   |    |
----------------------------------------------------
-3 |   |   |   |e1 |   |   |   |   |   |   |    |
----------------------------------------------------
-4 |   |   |   |   |e12|   |   |   |   |   |    |
----------------------------------------------------
-5 |   |   |   |   |   |e2 |   |   |   |   |    |
----------------------------------------------------
-6 |   |   |   |   |   |   |e1 |   |   |   |    |
----------------------------------------------------
-7 |   |   |   |   |   |   |   |e12|   |   |    |
----------------------------------------------------
-8 |   |   |   |   |   |   |   |   |e2 |   |    |
----------------------------------------------------
-9 |   |   |   |   |   |   |   |   |   |e1 |    |
----------------------------------------------------
-10|   |   |   |   |   |   |   |   |   |   |e12 |
----------------------------------------------------
-11|   |   |   |   |   |   |   |   |   |   |    |e2
-
 """
 
 # Standard imports
@@ -94,15 +49,15 @@
 
 # Local imports
 from SPI2py.models.physics.distributed.mesh import generate_mesh
-from SPI2py.models.physics.distributed.assembly import assemble_global_stiffness
-
+from SPI2py.models.physics.distributed.assembly import assemble_sparse_global_stiffness
 
 w = 1  # x
 h = 1  # y
 d = 1  # Z
 
+# Assemble the global stiffness matrix and convert it to a dense matrix
 nodes_2e, elements_2e, _, _, _, _, _, _, _ = generate_mesh(0, 2*w, 0, h, 0, d, element_size=1.0)
-Ke_flat, elem_indices, rows_flat, cols_flat, n_nodes, n_elem = assemble_global_stiffness(nodes_2e, elements_2e, base_k=1.0)
+Ke_flat, elem_indices, rows_flat, cols_flat, n_nodes, n_elem = assemble_sparse_global_stiffness(nodes_2e, elements_2e, base_k=1.0)
 indices = jnp.stack([rows_flat, cols_flat], axis=-1)
 K_global = K_pf = BCOO((Ke_flat, indices), shape=(n_nodes, n_nodes))
 K_global_dense = K_global.todense()
@@ -133,9 +88,6 @@
                          [H, G, F, E, D, C, B, A]])
 
 
-# K_global_expected = jnp.zeros((12, 12))
-# K_global_expected[]
-
 
 # Confirm that all entries that should be zeros are zeros
 
@@ -190,3 +142,122 @@
 assert jnp.all(K_global_dense[non_zeros_i, non_zeros_j] != 0.0)
 
 
+
+
+# Now spot check some values; K_e1 == K_e2 == K_ex
+
+# K[0,0] == K_e1[0,0]
+assert jnp.isclose(K_global_dense[0, 0], k_ex[0, 0])
+
+# K[1,1] == K_e1[1,1] + K_e2[0,0]
+assert jnp.isclose(K_global_dense[1, 1], k_ex[1, 1] + k_ex[0, 0])
+
+# K[0,9] == K_e1[0,7] and K[9,0] == K_e2[7,0]
+assert jnp.isclose(K_global_dense[0, 9], k_ex[0, 7])
+
+# K[11,11] == K_e2[6,6]
+assert jnp.isclose(K_global_dense[11, 11], k_ex[6, 6])
+
+
+
+
+
+
+
+
+
+
+
+
+
+# # Verify
+# cond_00_00 = 1
+# cond_00_01 = 1
+# cond_00_02 = 1
+# cond_00_03 = 1
+# cond_00_04 = 1
+# cond_00_05 = 1
+# cond_00_06 = 1
+# cond_00_07 = 1
+# cond_00_08 = 1
+# cond_00_09 = 1
+# cond_00_10 = 1
+# cond_00_11 = 1
+#
+# cond_01_00 = 1
+# cond_01_01 = 1
+# cond_01_02 = 1
+# cond_01_03 = 1
+# cond_01_04 = 1
+# cond_01_05 = 1
+# cond_01_06 = 1
+# cond_01_07 = 1
+# cond_01_08 = 1
+# cond_01_09 = 1
+# cond_01_10 = 1
+# cond_01_11 = 1
+#
+# cond_02_00 = 1
+# cond_02_01 = 1
+# cond_02_02 = 1
+# cond_02_03 = 1
+# cond_02_04 = 1
+# cond_02_05 = 1
+# cond_02_06 = 1
+# cond_02_07 = 1
+# cond_02_08 = 1
+# cond_02_09 = 1
+# cond_02_10 = 1
+# cond_02_11 = 1
+#
+# cond_03_00 = 1
+# cond_03_01 = 1
+# cond_03_02 = 1
+# cond_03_03 = 1
+# cond_03_04 = 1
+# cond_03_05 = 1
+# cond_03_06 = 1
+# cond_03_07 = 1
+# cond_03_08 = 1
+# cond_03_09 = 1
+# cond_03_10 = 1
+# cond_03_11 = 1
+#
+# cond_03_00 = 1
+# cond_03_01 = 1
+# cond_03_02 = 1
+# cond_03_03 = 1
+# cond_03_04 = 1
+# cond_03_05 = 1
+# cond_03_06 = 1
+# cond_03_07 = 1
+# cond_03_08 = 1
+# cond_03_09 = 1
+# cond_03_10 = 1
+# cond_03_11 = 1
+#
+# cond_04_00 = 1
+# cond_04_01 = 1
+# cond_04_02 = 1
+# cond_04_03 = 1
+# cond_04_04 = 1
+# cond_04_05 = 1
+# cond_04_06 = 1
+# cond_04_07 = 1
+# cond_04_08 = 1
+# cond_04_09 = 1
+# cond_04_10 = 1
+# cond_04_11 = 1
+#
+# cond_00_00 = 1
+# cond_00_01 = 1
+# cond_00_02 = 1
+# cond_00_03 = 1
+# cond_00_04 = 1
+# cond_00_05 = 1
+# cond_00_06 = 1
+# cond_00_07 = 1
+# cond_00_08 = 1
+# cond_00_09 = 1
+# cond_00_10 = 1
+# cond_00_11 = 1

examples/Scratch/scratch_test.py renamed to examples/Scratch/scratch_t.py

@@ -1,11 +1,13 @@
-import jax.numpy as jnp
+# Standard imports
 import matplotlib.pyplot as plt
-from SPI2py.models.physics.distributed.mesh import generate_mesh, find_active_nodes, find_face_nodes
-from SPI2py.models.physics.distributed.element import assemble_local_stiffness_matrix_scalar
+import jax.numpy as jnp
+from jax.experimental.sparse import BCOO
+
+# Local imports
+from SPI2py.models.physics.distributed.mesh import generate_mesh, find_face_nodes
 from SPI2py.models.physics.distributed.assembly import assemble_base_global_system_partition, apply_bc_partition
-from SPI2py.models.physics.distributed.quadrature import shape_functions
 from SPI2py.models.physics.distributed.solver import solve_system_partition
-
+from SPI2py.models.physics.distributed.assembly import assemble_sparse_global_stiffness
 
 
 """
@@ -26,19 +28,29 @@
 
 element_size = 0.25
 
+# Generate the mesh
 nodes, elements, centers, nx, ny, nz, lx, ly, lz = generate_mesh(x_min, x_max, y_min, y_max, z_min, z_max,
                                                                  element_size=element_size)
 
+
+
 densities = jnp.ones(elements.shape[0])  # uniform density
 heat_loads = jnp.zeros(elements.shape[0])  # no heat generation
 
 # Find active nodes and face nodes
-dirichlet_nodes = find_face_nodes(nodes, jnp.array([0.0, 0.0, -1.0]))
-robin_nodes = find_face_nodes(nodes, jnp.array([0.0, 0.0, 1.0]))
+top_normal = jnp.array([0, 1, 0])
+bottom_normal = jnp.array([0, -1, 0])
+
+dirichlet_nodes = find_face_nodes(nodes, bottom_normal)
+robin_nodes = find_face_nodes(nodes, top_normal)
+
 robin_area = element_size * element_size
 
 dirichlet_temperature = T_fixed * jnp.ones(dirichlet_nodes.shape[0])
 
+
+
+
 # Initialize global stiffness matrix and force vector
 K_base, f_base, elem_indices = assemble_base_global_system_partition(nodes, elements,
                                                                      base_k=k,

src/SPI2py/models/physics/distributed/assembly.py

@@ -7,7 +7,7 @@
 
 
 # @jit
-def assemble_global_stiffness(nodes, elements, base_k):
+def assemble_sparse_global_stiffness(nodes, elements, base_k):
     """
         Assemble the global base stiffness matrix.
         Indices are included for mapping element densities to each of their local stiffness matrix contributions.
@@ -60,6 +60,10 @@ def assemble_global_stiffness(nodes, elements, base_k):
     return Ke_flat, elem_indices, rows_flat, cols_flat, n_nodes, n_elem
 
 
+def partition_sparse_global_system():
+    pass
+
+
 # @jit
 def assemble_base_global_system_partition(nodes, elements,
                                           base_k,
@@ -70,8 +74,8 @@ def assemble_base_global_system_partition(nodes, elements,
     idx_p = d_nodes
     idx_f = jnp.setdiff1d(idx, idx_p)
 
-    Ke_flat, elem_indices_full, rows_flat, cols_flat, n_nodes, n_elem = assemble_global_stiffness(nodes, elements,
-                                                                                                  base_k)
+    Ke_flat, elem_indices_full, rows_flat, cols_flat, n_nodes, n_elem = assemble_sparse_global_stiffness(nodes, elements,
+                                                                                                         base_k)
     indices = jnp.stack([rows_flat, cols_flat], axis=-1)  # shape: (nnz, 2)
 
     # Partition the contributions into four blocks based on free (idx_f) and prescribed (idx_p) DOF indices.
@@ -134,8 +138,8 @@ def remap_local(global_inds, idx_array):
 
     # Apply boundary conditions (both Robin and Dirichlet).
     K_base, f_base = apply_bc_partition(K_base, f_base,
-                                      r_nodes, r_h, r_T_inf, r_area,
-                                      idx_f, idx_p)
+                                        r_nodes, r_h, r_T_inf, r_area,
+                                        idx_f, idx_p)
 
     return K_base, f_base, elem_indices
 
@@ -195,9 +199,9 @@ def update_global_stiffness_partition(K_base, f_base,
                                       element_densities, element_heat_loads):
 
     # Unpack the partitioned stiffness matrix, load vector, and their indices.
-    K_ff, K_fp, K_pf, K_pp         = K_base
-    f_f, f_p                       = f_base
-    idx_f, idx_p                   = idx_vector
+    K_ff, K_fp, K_pf, K_pp     = K_base
+    f_f, f_p                   = f_base
+    idx_f, idx_p               = idx_vector
     ei_ff, ei_fp, ei_pf, ei_pp = ei_matrix
 
     # Flatten the pseudo-densities and heat loads to ensure they are 1D.
@@ -251,7 +255,7 @@ def assemble_base_global_system_penalty(nodes, elements, base_k,
                                         r_nodes, r_h, r_T_inf, r_area,
                                         d_nodes, d_T):
 
-    Ke_flat, elem_indices, rows_flat, cols_flat, n_nodes, n_elem = assemble_global_stiffness(nodes, elements, base_k)
+    Ke_flat, elem_indices, rows_flat, cols_flat, n_nodes, n_elem = assemble_sparse_global_stiffness(nodes, elements, base_k)
 
     # Stack rows and cols to form an index array of shape (n_elem*64, 2).
     indices = jnp.stack([rows_flat, cols_flat], axis=-1)