code: """
  Model classes
  
  """
  
  import numpy as np
  import sys
  import inspect  # Used for storing the input
  from .aquifer import Aquifer
  from .aquifer_parameters import param_maq, param_3d
  from .constant import ConstantStar
  from .util import PlotTim
  import multiprocessing as mp
  
  __all__ = ['Model', 'ModelMaq', 'Model3D']
  
  class Model(PlotTim):
      """
      Model Class to create a model object consisting of an arbitrary
      sequence of aquifer layers and leaky layers.
      Use ModelMaq for regular sequence of aquifers and leaky layers.
      Use Model3D for multi-layer model of a single aquifer
      
      Parameters
      ----------
      kaq : array
          hydraulic conductivity of each aquifer from the top down
      z : array
          elevation tops and bottoms of all layers
          layers may have zero thickness
      c : array
          resistance between two consecutive aquifer layers
          if ltype[0]='a': length is number of aquifers - 1
          if ltype[0]='l': length is number of aquifers
      npor : array
          porosity of all layers from the top down
      ltype : array of characters
          array indicating for each layer whether it is
          'a' aquifer layer
          'l' leaky layer
      
      """
      
      def __init__(self, kaq, c, z, npor, ltype, f2py=False):
          # All input variables are numpy arrays
          # That should be checked outside this function
          self.elementlist = []
          self.elementdict = {}  # only elements that have a label
          self.aq = Aquifer(self, kaq, c, z, npor, ltype)
          self.modelname = 'ml'  # Used for writing out input
          self.f2py = False
          if f2py:
              try:
                  from .src import besselaesnew
                  self.f2py = True
              except:
                  print('FORTRAN extension not found while f2py=True')
                  print('Using Numba instead')
  
      def initialize(self):
          # remove inhomogeneity elements (they are added again)
          self.elementlist = [e for e in self.elementlist if not e.inhomelement]
          self.aq.initialize()
          for e in self.elementlist:
              e.initialize()
  
      def add_element(self, e):
          self.elementlist.append(e)
          if e.label is not None: self.elementdict[e.label] = e
  
      def remove_element(self, e):
          """Remove element `e` from model
          """
          
          if e.label is not None: self.elementdict.pop(e.label)
          self.elementlist.remove(e)
  
      def storeinput(self, frame):
          self.inputargs, _, _, self.inputvalues = inspect.getargvalues(frame)
  
      def potential(self, x, y, aq=None):
          if aq is None: aq = self.aq.find_aquifer_data(x, y)
          pot = np.zeros(aq.naq)
          for e in aq.elementlist:
              pot += e.potential(x, y, aq)
          rv = np.sum(pot * aq.eigvec, 1)
          if aq.ltype[0] == 'l':
              # potential for head above leaky layer
              rv += aq.constantstar.potstar
          return rv
  
      def disvec(self, x, y, aq=None):
          """Discharge vector at `x`, `y`
          
          Returns
          -------
          
          qxqy : array size (2, naq)
              first row is Qx in each aquifer layer, second row is Qy
          """
          
          if aq is None: aq = self.aq.find_aquifer_data(x, y)
          rv = np.zeros((2, aq.naq))
          for e in aq.elementlist:
              rv += e.disvec(x, y, aq)
          rv = np.sum(rv[:, np.newaxis, :] * aq.eigvec, 2)
          return rv
      
      def qztop(self, x, y, aq=None):
          if aq is None: aq = self.aq.find_aquifer_data(x, y)
          rv = 0.0
          if aq.ltype[0] == 'a':  # otherwise recharge cannot be added
              for e in aq.elementlist:
                  rv += e.qztop(x, y)
          return rv
  
      def head(self, x, y, layers=None, aq=None):
          """Head at `x`, `y`
          
          Returns
          -------
          
          h : array length `naq` or `len(layers)`
              head in all `layers` (if not `None`), 
              or all layers of aquifer (otherwise)
          """
          
          if aq is None: aq = self.aq.find_aquifer_data(x, y)
          rv = self.potential(x, y, aq) / aq.T
          if layers is None:
              return rv
          else:
              return rv[layers]
  
      def headgrid(self, xg, yg, layers=None, printrow=False):
          """Grid of heads
          
          Parameters
          ----------
          xg : array
              x values of grid
          yg : array
              y values of grid
          layers : integer, list or array, optional
              layers for which grid is returned
          printrow : boolean, optional
              prints dot to screen for each row of grid if set to `True`
          
          Returns
          -------
          h : array size `nlayers, ny, nx`
          
          See also
          --------
          
          :func:`~timml.model.Model.headgrid2`
  
          """
          
          nx, ny = len(xg), len(yg)
          if layers is None:
              Nlayers = self.aq.find_aquifer_data(xg[0], yg[0]).naq
          else:
              Nlayers = len(np.atleast_1d(layers))
          h = np.empty((Nlayers, ny, nx))
          for j in range(ny):
              if printrow:
                  print('.', end='', flush=True)
              for i in range(nx):
                  h[:, j, i] = self.head(xg[i], yg[j], layers)
          if printrow:
              print('', flush=True)
          return h
  
      def headgrid2(self, x1, x2, nx, y1, y2, ny, layers=None, printrow=False):
          """Grid of heads
          
          Parameters
          ----------
          x1, x2, nx : 
              x values are generated as linspace(x1, x2, nx)
          y1, y2, ny : 
              y values are generated as linspace(y1, y2, ny)
          layers : integer, list or array, optional
              layers for which grid is returned
          printrow : boolean, optional
              prints dot to screen for each row of grid if set to `True`
          
          Returns
          -------
          h : array size `nlayers, ny, nx`
          
          See also
          --------
          
          :func:`~timml.model.Model.headgrid`
          
          """
          
          xg, yg = np.linspace(x1, x2, nx), np.linspace(y1, y2, ny)
          return self.headgrid(xg, yg, layers=layers, printrow=printrow)
  
      def headalongline(self, x, y, layers=None):
          """Head along line or curve
          
          Parameters
          ----------
          x : array
              x values of line
          y : array
              y values of line
          layers : integer, list or array, optional
              layers for which grid is returned
          
          Returns
          -------
          h : array size `nlayers, nx`
  
          """
          
          xg, yg = np.atleast_1d(x), np.atleast_1d(y)
          if layers is None:
              Nlayers = self.aq.find_aquifer_data(xg[0], yg[0]).naq
          else:
              Nlayers = len(np.atleast_1d(layers))
          nx = len(xg)
          if len(yg) == 1:
              yg = yg * np.ones(nx)
          h = np.zeros((Nlayers, nx))
          for i in range(nx):
              h[:, i] = self.head(xg[i], yg[i], layers)
          return h
      
      def disvecalongline(self, x, y, layers=None):
          '''Returns Qx[Nlayers,len(x)], Qy[Nlayers,len(x)]
          Assumes same number of layers for each x and y
          layers may be None or list of layers for which head is computed'''
          xg, yg = np.atleast_1d(x), np.atleast_1d(y)
          if layers is None:
              nlayers = self.aq.find_aquifer_data(xg[0], yg[0]).naq
          else:
              nlayers = len(np.atleast_1d(layers))
          nx = len(xg)
          if len(yg) == 1:
              yg = yg * np.ones(nx)
          Qx = np.zeros((nlayers, nx))
          Qy = np.zeros((nlayers, nx))
          for i in range(nx):
              Qx[:, i], Qy[:, 1] = self.disvec(xg[i], yg[i], layers)
          return Qx, Qy
      
  #    def disvec_direction(self, s, x1, y1, cdirection):
  #        pass
  #    
  #    def discharge_across_line(self, x1, y1, x2,  y2, layers=None):
  #        if layers is None:
  #            nlayers = self.aq.find_aquifer_data(x1, y1).naq
  #        else:
  #            nlayers = len(np.atleast_1d(layers))
  #        z1 = x1 + y1 * 1j
  #        z2 = x2 + y2 * 1j
  #        normvec = (z2 - z1) / np.abs(z2 - z1) * np.exp(-np.pi * 1j / 2)
  #        disvec = self.disvec(xg[i], yg[i], layers)
      
      def velocity(self, x, y, z):
          return self.velocomp(x, y, z)
      
      def velocomp(self, x, y, z, aq=None, layer_ltype=None):
          if aq is None: aq = self.aq.find_aquifer_data(x, y)
          assert z <= aq.z[0] and z >= aq.z[-1], "z value not inside aquifer"
          if layer_ltype is None:
              layer, ltype, dummy = aq.findlayer(z)
          else:
              layer, ltype = layer_ltype
          h = self.head(x, y, aq=aq)
          # qz between aquifer layers
          qzlayer = np.zeros(aq.naq + 1)
          qzlayer[1:-1] = (h[1:] - h[:-1]) / aq.c[1:]
          if aq.ltype[0] == 'l':
              qzlayer[0] = (h[0] - aq.hstar) / aq.c[0]
          if ltype == 'l':
              vz = qzlayer[layer] / aq.nporll[layer]
              vx = 0
              vy = 0
          else:
              qzbot = qzlayer[layer + 1]
              qztop = qzlayer[layer]
              if layer == 0:
                  qztop += self.qztop(x, y)   
              vz = (qzbot + (z - aq.zaqbot[layer]) / aq.Haq[layer] * \
                   (qztop - qzbot)) / aq.nporaq[layer]        
              qx, qy = self.disvec(x, y, aq=aq)
              vx = qx[layer] / (aq.Haq[layer] * aq.nporaq[layer])
              vy = qy[layer] / (aq.Haq[layer] * aq.nporaq[layer])
          return np.array([vx, vy, vz])
                 
      def solve(self, printmat=0, sendback=0, silent=False):
          '''Compute solution'''
          # Initialize elements
          self.initialize()
          # Compute number of equations
          self.neq = np.sum([e.nunknowns for e in self.elementlist])
          if self.neq == 0: return
          if silent is False:
              print('Number of elements, Number of equations:', len(
                  self.elementlist), ',', self.neq)
          if self.neq == 0:
              if silent is False: print('No unknowns. Solution complete')
              return
          mat = np.empty((self.neq, self.neq))
          rhs = np.empty(self.neq)
          ieq = 0
          for e in self.elementlist:
              if e.nunknowns > 0:
                  mat[ieq:ieq + e.nunknowns, :], rhs[ieq:ieq + e.nunknowns] = \
                  e.equation()
                  ieq += e.nunknowns
              if silent is False:
                  print('.', end='', flush=True)
          if printmat:
              return mat, rhs
          sol = np.linalg.solve(mat, rhs)
          icount = 0
          for e in self.elementlist:
              if e.nunknowns > 0:
                  e.setparams(sol[icount:icount + e.nunknowns])
                  icount += e.nunknowns
          if silent is False:
              print()  # needed cause the dots are printed
              print('solution complete')
          elif (silent == 'dot') or (silent == '.'):
              print('.', end='', flush=True)
          if sendback:
              return sol
          return
  
      def solve_mp(self, nproc=4, printmat=0, sendback=0, silent=False):
          '''Compute solution, multiprocessing implementation.
          Note: estimated speedup approximately by factor of
          number of physical cores. Virtual cores do not improve
          calculation time.'''
          # Initialize elements
          self.initialize()
          # Compute number of equations
          self.neq = np.sum([e.nunknowns for e in self.elementlist])
          if self.neq == 0: return
          if silent is False:
              print('Number of elements, Number of equations:', len(
                  self.elementlist), ',', self.neq)
          if self.neq == 0:
              if silent is False: print('No unknowns. Solution complete')
              return
          mat = np.empty((self.neq, self.neq))
          rhs = np.empty(self.neq)
  
          # start multiprocessing
          if nproc is None:
              nproc = mp.cpu_count() - 1  # make no. of processes equal to 1 less than no. of cores
          elif nproc > mp.cpu_count():
              print("Given 'nproc' larger than no. of cores on machine. Setting 'nproc' to {}.".format(mp.cpu_count()))
              nproc = mp.cpu_count()
  
          pool = mp.Pool(processes=nproc)
          results = []
          for e in self.elementlist:
              if e.nunknowns > 0:
                  results.append(pool.apply_async(e.equation))
              if silent is False:
                  print('.', end='', flush=True)
  
          pool.close()
          pool.join()
  
          mat = np.empty((self.neq, self.neq))
          rhs = np.zeros(self.neq)
  
          ieq = 0
  
          for p in results:
              imat, irhs = p.get()
              mat[ieq:ieq + imat.shape[0], :] = imat
              rhs[ieq:ieq + irhs.shape[0]] = irhs
              ieq += imat.shape[0]
  
          # end multiprocessing
  
          if printmat:
              return mat, rhs
          sol = np.linalg.solve(mat, rhs)
          icount = 0
          for e in self.elementlist:
              if e.nunknowns > 0:
                  e.setparams(sol[icount:icount + e.nunknowns])
                  icount += e.nunknowns
          if silent is False:
              print()  # needed cause the dots are printed
              print('solution complete')
          elif (silent == 'dot') or (silent == '.'):
              print('.', end='', flush=True)
          if sendback:
              return sol
          return
      
      def write(self):
          rv = self.modelname + ' = ' + self.name + '(\n'
          for key in self.inputargs[1:]:  # The first argument (self) is ignored
              if isinstance(self.inputvalues[key], np.ndarray):
                  rv += key + ' = ' + np.array2string(self.inputvalues[key], 
                                                      separator=',') + ',\n'
              elif isinstance(self.inputvalues[key],str):                
                  rv += key + " = '" + self.inputvalues[key] + "',\n"
              else:
                  rv += key + ' = ' + str(self.inputvalues[key]) + ',\n'
          rv += ')\n'
          return rv
      
      def writemodel(self, fname):
          self.initialize()  # So that the model can be written without solving first
          f = open(fname, 'w')
          f.write('from timml import *\n')
          f.write(self.write())
          for e in self.elementlist:
              f.write(e.write())
          f.close()
          
  class ModelMaq(Model):
      """
      Create a Model object by specifying a mult-aquifer sequence of
      aquifer-leakylayer-aquifer-leakylayer-aquifer etc
      
      Parameters
      ----------
      kaq : float, array or list
          Hydraulic conductivity of each aquifer from the top down.
          If float, hydraulic conductivity is the same in all aquifers.
      z : array or list
          Elevation of tops and bottoms of the aquifers from the top down.
          Leaky layers may have zero thickness.
             * if topboundary='conf': length is 2 * number of aquifers
             * if topboundary='semi': length is 2 * number of aquifers + 1 
               as top of leaky layer on top of systems needs to be specified
      c : float, array or list
          Resistance of leaky layers from the top down.
             * if float, resistance is the same for all leaky layers
             * if topboundary='conf': length is number of aquifers - 1
             * if topboundary='semi': length is number of aquifers
      npor : float, array or list
          Porosity of all aquifers and leaky layers from the top down.
             * if float, porosity is the same for all layers
             * if topboundary='conf': length is 2 * number of aquifers - 1
             * if topboundary='semi': length is 2 * number of aquifers
      topboundary : string, 'conf' or 'semi' (default is 'conf')
          Indicates whether the topboundary is confined ('conf') or
          semi-confined ('semi').
      hstar : float or None (default is None)
          Head value above semi-confining top, only read if topboundary='semi'.
  
      Examples
      --------
      >>> ml = ModelMaq(kaq=[10, 20], z=[20, 12, 10, 0], c=1000)
      
      """
      
      def __init__(self, kaq=1, z=[1, 0], c=[], npor=0.3, topboundary='conf',
                   hstar=None, f2py=False):
          self.storeinput(inspect.currentframe())
          kaq, c, npor, ltype = param_maq(kaq, z, c, npor, topboundary)
          Model.__init__(self, kaq, c, z, npor, ltype, f2py)
          self.name = 'ModelMaq'
          if self.aq.ltype[0] == 'l':
              ConstantStar(self, hstar, aq=self.aq)
              
  class Model3D(Model):
      """
      Model3D Class to create a multi-layer model object consisting of
      many aquifer layers. The resistance between the layers is computed
      from the vertical hydraulic conductivity of the layers.
      
      Parameters
      ----------
      kaq : float, array or list
          hydraulic conductivity of each layer from the top down
          if float, hydraulic conductivity is the same in all aquifers
      z : array or list
          elevation of top of system followed by bottoms of all layers
          from the top down
          bottom of layer is automatically equal to top of layer below it
          length is number of aquifer layers + 1
      kzoverkh : float
          vertical anisotropy ratio vertical k divided by horizontal k
          if float, value is the same for all layers
          length is number of layers
      npor : float, array or list
          porosity of all aquifer layers
          from the top down
          if float, porosity is the same for all layers
          if topboundary='conf': length is number of layers
          if topboundary='semi': length is number of layers + 1
      topboundary : string, 'conf' or 'semi' (default is 'conf')
          indicating whether the top is confined ('conf') or
          semi-confined ('semi')
      topres : float
          resistance of top semi-confining layer (read if topboundary='semi')
      topthick: float
          thickness of top semi-confining layer (read if topboundary='semi')
      hstar : float or None (default is None)
          head value above semi-confining top (read if topboundary='semi')
  
      Examples
      --------
      >>> ml = Model3D(kaq=10, z=np.arange(20, -1, -2), kzoverkh=0.1)
      
      """
      
      def __init__(self, kaq=1, z=[1, 0], kzoverkh=1, npor=0.3,
                   topboundary='conf', topres=0, topthick=0, hstar=0,
                   f2py=False):
          '''Model3D
          for semi-confined aquifers, set top equal to 'semi' and provide
          topres: resistance of top
          tophick: thickness of top
          hstar: head above top'''
          self.storeinput(inspect.currentframe())
          kaq, c, npor, ltype = param_3d(kaq, z, kzoverkh, npor, topboundary,
                                         topres)
          if topboundary == 'semi':
              z = np.hstack((z[0] + topthick, z))
          Model.__init__(self, kaq, c, z, npor, ltype, f2py)
          self.name = 'Model3D'
          if self.aq.ltype[0] == 'l':
              ConstantStar(self, hstar, aq=self.aq)
  
  
apis: numpy.sum numpy.ones numpy.linalg.solve numpy.empty numpy.zeros numpy.array2string numpy.atleast_1d numpy.hstack numpy.array numpy.linspace