start modifying disvec in Model, Element and Well classes

- TODO: check the adjustments for the methods currently not used by the Well class and example notebook

ttim/element.py

@@ -120,10 +120,10 @@ def unitpotentialone(self, x, y, jtime, aq=None):
         return np.sum(self.potinfone(x, y, jtime, aq), 0)
 
     @abstractmethod
-    def disvecinf(self, x, y, aq=None):
-        """Returns 2 complex arrays of size (nparam, naq, npval)."""
+    def disvecinf(self, x, y, t_int, aq=None):
+        """Returns 2 complex arrays of size (nparam, naq, nppar)."""
 
-    def disvec(self, x, y, aq=None):
+    def disvecall(self, x, y, aq=None):
         """Returns 2 complex arrays of size (ngvbc, naq, npval)."""
         if aq is None:
             aq = self.model.aq.find_aquifer_data(x, y)
@@ -132,15 +132,25 @@ def disvec(self, x, y, aq=None):
             self.parameters[:, :, np.newaxis, :] * qy, 1
         )
 
-    def unitdisvec(self, x, y, aq=None):
-        """Returns 2 complex arrays of size (naq, npval).
+    def disvec(self, x, y, t_int, aq=None):
+        """Returns 2 complex arrays of size (ngvbc, naq, nppar)."""
+        if aq is None:
+            aq = self.model.aq.find_aquifer_data(x, y)
+        qx, qy = self.disvecinf(x, y, t_int, aq)
+        return (
+            np.sum(self.parameters[t_int] * qx, 1),
+            np.sum(self.parameters[t_int] * qy, 1),
+        )
+
+    def unitdisvec(self, x, y, t_int, aq=None):
+        """Returns 2 complex arrays of size (naq, nppar).
 
         Can be more efficient for given elements.
         """
         if aq is None:
             aq = self.model.aq.find_aquifer_data(x, y)
-        qx, qy = self.disvecinf(x, y, aq)
-        return np.sum(qx, 0), np.sum(qy, 0)
+        qx, qy = self.disvecinf(x, y, t_int, aq)
+        return np.sum(qx, axis=0), np.sum(qy, axis=0)
 
     # Functions used to build equations
     def potinflayers(self, x, y, t_int, layers=0, aq=None):
@@ -178,44 +188,44 @@ def unitpotentiallayers(self, x, y, layers=0, aq=None):
         phi = np.sum(pot[np.newaxis, :, :] * aq.eigvec, 1)
         return phi[layers, :]
 
-    def disvecinflayers(self, x, y, layers=0, aq=None):
+    def disvecinflayers(self, x, y, t_int, layers=0, aq=None):
         """Layers can be scalar, list, or array.
 
         returns 2 arrays of size (len(layers),nparam,npval) only used in building
         equations
         """
         if aq is None:
             aq = self.model.aq.find_aquifer_data(x, y)
-        qx, qy = self.disvecinf(x, y, aq)
-        rvx = np.sum(qx[:, np.newaxis, :, :] * aq.eigvec, 2)
-        rvy = np.sum(qy[:, np.newaxis, :, :] * aq.eigvec, 2)
+        qx, qy = self.disvecinf(x, y, t_int, aq)
+        rvx = qx * aq.eigvec[t_int]
+        rvy = qy * aq.eigvec[t_int]
         # first axis needs to be the number of layers
         rvx = rvx.swapaxes(0, 1)
         rvy = rvy.swapaxes(0, 1)
         return rvx[layers, :], rvy[layers, :]
 
-    def disveclayers(self, x, y, layers=0, aq=None):
-        """Returns 2 complex array of size (ngvbc, len(layers), npval).
+    def disveclayers(self, x, y, t_int, layers=0, aq=None):
+        """Returns 2 complex array of size (ngvbc, len(layers), nppar).
 
         Only used in building equations.
         """
         if aq is None:
             aq = self.model.aq.find_aquifer_data(x, y)
-        qx, qy = self.disvec(x, y, aq)
-        rvx = np.sum(qx[:, np.newaxis, :, :] * aq.eigvec, 2)
-        rvy = np.sum(qy[:, np.newaxis, :, :] * aq.eigvec, 2)
+        qx, qy = self.disvec(x, y, t_int, aq)
+        rvx = qx * aq.eigvec[t_int]
+        rvy = qy * aq.eigvec[t_int]
         return rvx[:, layers, :], rvy[:, layers, :]
 
-    def unitdisveclayers(self, x, y, layers=0, aq=None):
-        """Returns complex array of size (len(layers), npval).
+    def unitdisveclayers(self, x, y, t_int, layers=0, aq=None):
+        """Returns complex array of size (len(layers), nppar).
 
         Only used in building equations.
         """
         if aq is None:
             aq = self.model.aq.find_aquifer_data(x, y)
-        qx, qy = self.unitdisvec(x, y, aq)
-        rvx = np.sum(qx[np.newaxis, :, :] * aq.eigvec, 1)
-        rvy = np.sum(qy[np.newaxis, :, :] * aq.eigvec, 1)
+        qx, qy = self.unitdisvec(x, y, t_int, aq)
+        rvx = np.sum(qx * aq.eigvec[t_int], axis=1)
+        rvy = np.sum(qy * aq.eigvec[t_int], axis=1)
         return rvx[layers, :], rvy[layers, :]
 
     def discharge(self, t, derivative=0):

ttim/model.py

@@ -316,6 +316,64 @@ def disvec(self, x, y, t, layers=None, aq=None, derivative=0):
         )
         return rvx, rvy
 
+    def disvec(self, x, y, time, layers=None, aq=None, derivative=0):
+        """Compute discharge vectgor.
+
+        Returns qx[naq, ntimes], qy[naq, ntimes] if layers=None, otherwise
+        qx[len(layers,Ntimes)],qy[len(layers, ntimes)].
+
+        t must be ordered.
+        """
+        if aq is None:
+            aq = self.aq.find_aquifer_data(x, y)
+        if layers is None:
+            layers = range(aq.naq)
+        nlayers = len(layers)
+        t_int = np.floor(np.log10(time)).astype(int)
+        time = np.atleast_1d(time) - self.tstart
+        disx = np.zeros((self.ngvbc, aq.naq, self.nppar), dtype=complex)
+        disy = np.zeros((self.ngvbc, aq.naq, self.nppar), dtype=complex)
+        for i in range(self.ngbc):
+            qx, qy = self.gbclist[i].unitdisvec(x, y, t_int, aq)
+            disx[i, :] += qx
+            disy[i, :] += qy
+        for e in self.vzbclist:
+            qx, qy = e.disvec(x, y, t_int, aq)
+            disx += qx
+            disy += qy
+        if layers is None:
+            disx = disx * aq.eigvec[t_int]
+            disy = disy * aq.eigvec[t_int]
+        else:
+            disx = disx * aq.eigvec[t_int][layers, :]
+            disy = disy * aq.eigvec[t_int][layers, :]
+        if derivative > 0:
+            disx *= self.p**derivative
+            disy *= self.p**derivative
+        rvx = invlapcomp(
+            time,
+            disx,
+            self.nppar,
+            self.M,
+            self.tintervals[t_int],
+            self.enumber,
+            self.etstart,
+            self.ebc,
+            nlayers,
+        )
+        rvy = invlapcomp(
+            time,
+            disy,
+            self.nppar,
+            self.M,
+            self.tintervals[t_int],
+            self.enumber,
+            self.etstart,
+            self.ebc,
+            nlayers,
+        )
+        return rvx, rvy
+
     def head(self, x, y, t, layers=None, aq=None, derivative=0, neglect_steady=False):
         """Head at x, y, t where t can be multiple times.
 

ttim/well.py

@@ -127,16 +127,18 @@ def potinf(self, x, y, t_int, aq=None):
                     rv[:, i, :] = self.term[t_int][:, i, :] * pot
         return rv
 
-    def disvecinf(self, x, y, aq=None):
+    def disvecinfall(self, x, y, aq=None):
         """Can be called with only one x,y value."""
         if aq is None:
             aq = self.model.aq.find_aquifer_data(x, y)
-        qx = np.zeros((self.nparam, aq.naq, self.model.npval), "D")
-        qy = np.zeros((self.nparam, aq.naq, self.model.npval), "D")
+        qx = np.zeros((self.nparam, aq.naq, self.model.npval), dtype=complex)
+        qy = np.zeros((self.nparam, aq.naq, self.model.npval), dtype=complex)
         if aq == self.aq:
-            qr = np.zeros((self.nparam, aq.naq, self.model.nint, self.model.nppar), "D")
+            qr = np.zeros(
+                (self.nparam, aq.naq, self.model.nint, self.model.nppar), dtype=complex
+            )
             r = np.sqrt((x - self.xw) ** 2 + (y - self.yw) ** 2)
-            # pot = np.zeros(self.model.nppar, "D")
+            # pot = np.zeros(self.model.nppar, dtype=complex)
             if r < self.rw:
                 r = self.rw  # If at well, set to at radius
             for i in range(self.aq.naq):
@@ -152,6 +154,29 @@ def disvecinf(self, x, y, aq=None):
             qy[:] = qr * (y - self.yw) / r
         return qx, qy
 
+    def disvecinf(self, x, y, t_int, aq=None):
+        """Can be called with only one x, y value for log time interval t_int."""
+        if aq is None:
+            aq = self.model.aq.find_aquifer_data(x, y)
+        qx = np.zeros((self.nparam, aq.naq, self.model.nppar), dtype=complex)
+        qy = np.zeros((self.nparam, aq.naq, self.model.nppar), dtype=complex)
+        if aq == self.aq:
+            qr = np.zeros((self.nparam, aq.naq, self.model.nppar), dtype=complex)
+            r = np.sqrt((x - self.xw) ** 2 + (y - self.yw) ** 2)
+            # pot = np.zeros(self.model.nppar, dtype=complex)
+            if r < self.rw:
+                r = self.rw  # If at well, set to at radius
+            for i in range(self.aq.naq):
+                if r / abs(self.aq.lab[t_int][i, 0]) < self.rzero:
+                    qr[:, i, :] = (
+                        self.term[t_int][:, i, :]
+                        * kv(1, r / self.aq.lab[t_int][i, :])
+                        / self.aq.lab[t_int][i, :]
+                    )
+            qx[:] = qr * (x - self.xw) / r
+            qy[:] = qr * (y - self.yw) / r
+        return qx, qy
+
     def headinside(self, t, derivative=0):
         """Returns head inside the well for the layers that the well is screened in.