Improve appearance of output image to see the sink clearly (insulator is black)

diff --git a/resistor_solver.py b/resistor_solver.py
index 8529da9..55cf061 100755 (executable)
--- a/resistor_solver.py
+++ b/resistor_solver.py
@@ -82,52 +82,77 @@ A = scipy.sparse.coo_array(
   dtype = numpy.int32
 ).tocsr()
 
-# partition nodes into 3 groups:
-# 0: black + blue (electrode), 1: red (electrode), 2: white (conductor)
-# note: black pixels are part of electrode to give them a defined voltage
-g0 = (image < 2).reshape((n,))
+# partition nodes into groups:
+# 0: black (insulator)
+# 1: blue (sink)
+# 2: red (source)
+# 3: white (conductor)
+g0 = (image == 0).reshape((n,))
 n0 = numpy.sum(g0)
-g1 = (image == 2).reshape((n,))
+g1 = (image == 1).reshape((n,))
 n1 = numpy.sum(g1)
-g2 = (image == 3).reshape((n,))
+g2 = (image == 2).reshape((n,))
 n2 = numpy.sum(g2)
-print('n0', n0, 'n1', n1, 'n2', n2)
+g3 = (image == 3).reshape((n,))
+n3 = numpy.sum(g3)
+print('n0', n0, 'n1', n1, 'n2', n2, 'n3', n3)
 
-# partition A into 3 x 3 submatrices by groups, V and I into subvectors
+# partition A into 4 x 4 submatrices by groups, V and I into subvectors
 #A00 = A[g0, :][:, g0]
 #A01 = A[g0, :][:, g1]
 #A02 = A[g0, :][:, g2]
+#A03 = A[g0, :][:, g3]
 #A10 = A[g1, :][:, g0]
-A11 = A[g1, :][:, g1]
-A12 = A[g1, :][:, g2]
+#A11 = A[g1, :][:, g1]
+#A12 = A[g1, :][:, g2]
+#A13 = A[g1, :][:, g3]
 #A20 = A[g2, :][:, g0]
-A21 = A[g2, :][:, g1]
+#A21 = A[g2, :][:, g1]
 A22 = A[g2, :][:, g2]
+A23 = A[g2, :][:, g3]
+#A30 = A[g3, :][:, g0]
+#A31 = A[g3, :][:, g1]
+A32 = A[g3, :][:, g2]
+A33 = A[g3, :][:, g3]
 
 # solve:
-# A00 V0 + A01 V1 + A02 V2 = I0
-# A10 V0 + A11 V1 + A12 V2 = I1
-# A20 V0 + A21 V1 + A22 V2 = I2
-# where V0 = all 0, V1 = all 1, I2 = all 0 => unknowns are I0, I1, V2
-#A01_V1 = numpy.sum(A01.astype(numpy.int32), 1)
-A11_V1 = numpy.sum(A11.astype(numpy.int32), 1)
-A21_V1 = numpy.sum(A21.astype(numpy.int32), 1)
+# A00 V0 + A01 V1 + A02 V2 + A03 V3 = I0
+# A10 V0 + A11 V1 + A12 V2 + A13 V3 = I1
+# A20 V0 + A21 V1 + A22 V2 + A23 V3 = I2
+# A30 V0 + A31 V1 + A32 V2 + A33 V3 = I3
+# where V0 = all 0, V1 = all 0, V2 = all 1, I0 = all 0, I2 = all 0
+# => unknowns are I1, I2, V3
+#A02_V2 = numpy.sum(A02.astype(numpy.int32), 1)
+#A12_V2 = numpy.sum(A12.astype(numpy.int32), 1)
+A22_V2 = numpy.sum(A22.astype(numpy.int32), 1)
+A32_V2 = numpy.sum(A32.astype(numpy.int32), 1)
 print('solve')
-V2 = scipy.sparse.linalg.spsolve(
-  A22.astype(numpy.double),
-  -A21_V1.astype(numpy.double)
+V3 = scipy.sparse.linalg.spsolve(
+  A33.astype(numpy.double),
+  -A32_V2.astype(numpy.double)
 )
 print('solve done')
 
-# we only care about sum(I0) and sum(I1), and we know sum(I0) = -sum(I1)
-sum_I1 = numpy.sum(A11_V1, 0) + numpy.sum(A12, 0) @ V2
-print('R', 1. / sum_I1, 'squares')
+# we only care about sum(I1) and sum(I2), and we know sum(I1) = -sum(I2)
+sum_I2 = numpy.sum(A22_V2, 0) + numpy.sum(A23, 0) @ V3
+print('R', 1. / sum_I2, 'squares')
 
-V = numpy.zeros((n,), numpy.double)
-V[g1] = 1.
-V[g2] = V2
+#V = numpy.zeros((n,), numpy.double)
+#V[g2] = 1.
+#V[g3] = V3
+
+# construct image as follows:
+# insulator is solid black
+# sink is solid blue
+# source is solid red
+# conductor fades from blue to red with increasing voltage
+image = numpy.zeros((n, 3), numpy.double)
+image[g1, 2] = 1.
+image[g2, 0] = 1.
+image[g3, 2] = 1. - V3
+image[g3, 0] = V3
+image = image.reshape((ys, xs, 3))
 
-image = V.reshape((ys, xs))
 mask = image >= .04045
 image[~mask] /= 12.92
 image[mask] = ((image[mask] + .055) / 1.055) ** 2.4