From 6937bbf455ad40f4bf82be860949628425f87e8d Mon Sep 17 00:00:00 2001
From: Jan Petykiewicz <jan@mpxd.net>
Date: Mon, 18 Mar 2024 10:59:00 -0700
Subject: [PATCH] use slices

---
 nom-eme.py | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/nom-eme.py b/nom-eme.py
index d9c7c8b..33d42f6 100644
--- a/nom-eme.py
+++ b/nom-eme.py
@@ -1,6 +1,6 @@
 import scipy
 import numpy
-from numpy.typing import ArrayLike, NDarray
+from numpy.typing import ArrayLike, NDArray
 
 
 #from simphony.elements import Model
@@ -262,9 +262,12 @@ def connect_s(
     nC = nA + nB - 2
     assert numpy.array_equal(A.shape[:-2], B.shape[:-2])
 
-    denom = 1 - A[..., k, k] * B[..., l, l]
-    Anew = A + A[..., k, :] * B[..., l, l] * A[..., :, k] / denom
-    Bnew = A[..., k, :] * B[..., :, l] / denom
+    ll = slice(l, l + 1)
+    kk = slice(k, k + 1)
+
+    denom = 1 - A[..., kk, kk] * B[..., ll, ll]
+    Anew = A + A[..., kk, :] * B[..., ll, ll] * A[..., :, kk] / denom
+    Bnew = A[..., kk, :] * B[..., :, ll] / denom
     Anew = npy.delete(Anew, (k,), 1)
     Anew = npy.delete(Anew, (k,), 2)
     Bnew = npy.delete(Bnew, (l,), 1)