grid.py

   1 #-*- coding: utf-8 -*-
   2
   3 import numpy as np
   4 from scipy import weave
   5 from scipy.weave import blitz, inline, converters
   6 from const import timestep
   7 from sweep import native_sweep, sweep
   8
   9 def norm(lhs, rhs):
  10     if not (isinstance(lhs, np.ndarray)
  11             and isinstance(rhs, np.ndarray)):
  12         raise TypeError("arguments must be arrays")
  13
  14     if lhs.shape != rhs.shape:
  15         raise ValueError("arrays must be the same shape")
  16
  17     return np.max(np.abs(lhs - rhs))
  18
  19
  20 class ExplicitSolver(object):
  21     '''Represents equation solver with time stepper
  22     '''
  23
  24     def __init__(self, I, J, time, ht=None, ly=2e-5, lz=2e-5, lambda_=2e-6, C=3e8, solver='inline', debug=False):
  25
  26         if solver == 'python':
  27             self._step = self._python_step
  28         elif solver == 'blitz':
  29             self._step = self._blitz_step
  30         elif solver == 'inline':
  31             self._step = self._inline_step
  32         else:
  33             raise ValueError('incorrect solver: %s' % solver)
  34
  35         self.solve = self._explicit_solve
  36
  37         self.C = C
  38         self.time = time
  39         self.I, self.J, self.ly, self.lz, self.lambda_ = I, J, ly, lz, lambda_
  40         self.hy = np.double(ly / I)
  41         self.hz = np.double(lz / J)
  42         if ht is None:
  43             self.ht = timestep(min(self.hy, self.hz) / (C * np.sqrt(2)))
  44         else:
  45             self.ht = ht
  46         self.K = int(self.time / self.ht)
  47         self.u = np.zeros((3, I + 1, J + 1), 'd')
  48
  49         self._init_bc(0, 0)
  50         self._init_bc(1, 1)
  51         self._init_bc(2, 2)
  52         self._solved = False
  53
  54         if debug:
  55             self._debug_print()
  56
  57     def _debug_print(self):
  58         print "grid: %d x %d" % (self.I, self.J)
  59         print "time:ht:K = %g:%g:%d" % (self.time, self.ht, self.K)
  60         print
  61         print "ly:hy:I = %g:%g:%d" % (self.ly, self.hy, self.I)
  62         print "lz:hz:J = %g:%g:%d" % (self.lz, self.hz, self.J)
  63         print
  64         print "u size (in Kbytes):", ((self.u.size * self.u.itemsize) / 1024.)
  65         print "extrapolated size (in Mbytes):", (float(self.I * self.J * self.K * self.u.itemsize) / (1024**2))
  66         print
  67
  68     def _init_bc(self, k, time):
  69         for i in xrange(self.u.shape[1]):
  70             self.u[k, i, 0] = np.sin((np.pi * i) / self.I) * np.sin((2 * np.pi * self.C * self.ht * time) / self.lambda_)
  71
  72     def _python_step(self):
  73         g = (self.C**2 * self.ht**2) / self.hy**2
  74         d = (self.C**2 * self.ht**2) / self.hz**2
  75         self.u[2, 1:-1, 1:-1] = \
  76              g * (self.u[1, 2:, 1:-1] + self.u[1, 0:-2, 1:-1]) + \
  77              d * (self.u[1, 1:-1, 2:] + self.u[1, 1:-1, 0:-2]) + \
  78              2 * (1 - g - d) * self.u[1, 1:-1, 1:-1] - \
  79              self.u[0, 1:-1, 1:-1]
  80
  81     def _blitz_step(self):
  82         g = float((self.C**2 * self.ht**2) / self.hy**2)
  83         d = float((self.C**2 * self.ht**2) / self.hz**2)
  84         u = self.u
  85         expr = \
  86         "u[2, 1:-1, 1:-1] = " \
  87             "g * (u[1, 2:, 1:-1] + u[1, 0:-2, 1:-1]) + " \
  88              "d * (u[1, 1:-1, 2:] + u[1, 1:-1, 0:-2]) + " \
  89              "2 * (1 - g - d)  * u[1, 1:-1, 1:-1] - " \
  90              "u[0, 1:-1, 1:-1]"
  91         blitz(expr, check_size=0)
  92
  93     def _inline_step(self):
  94         g = float((self.C**2 * self.ht**2) / self.hy**2)
  95         d = float((self.C**2 * self.ht**2) / self.hz**2)
  96         u = self.u
  97         I, J, K = self.I, self.J, self.K
  98         code = """
  99         for (int i=1; i < I-1; ++i) {
 100             for (int j = 1; j < J-1; ++j) {
 101                 u(2, i, j) = g * (u(1, i+1, j) + u(1, i-1, j)) +
 102                              d * (u(1, i, j+1) + u(1, i, j-1)) +
 103                              2 * (1 - g - d) * u(1, i, j) -
 104                              u(0, i, j);
 105             }
 106         }
 107         """
 108         inline(code,
 109                ['u', 'g', 'd', 'I', 'J', 'K'],
 110                type_converters=converters.blitz,
 111                compiler='gcc')
 112
 113     def _explicit_solve(self):
 114         if not self._solved:
 115             for k in xrange(2, self.K):
 116                 self._step()
 117                 self.u[0] = self.u[1].copy()
 118                 self.u[1] = self.u[2].copy()
 119                 self._init_bc(2, k + 1)
 120             self._solved = True
 121         return self.u[2]
 122
 123
 124 class ImplicitSolver(object):
 125     '''Represents equation solver with implicit difference scheme
 126     '''
 127
 128     def __init__(self, I, J, time, ht=None, ly=2e-5, lz=2e-5, lambda_=2e-6, C=3e8, debug=False):
 129
 130         self.solve = self._implicit_solve
 131
 132         self.C = C
 133         self.time = time
 134         self.I, self.J, self.ly, self.lz, self.lambda_ = I, J, ly, lz, lambda_
 135         self.hy = np.double(ly / I)
 136         self.hz = np.double(lz / J)
 137         if ht is None:
 138             val = self.hy / C
 139             s = "%g" % val
 140             l = s.split('e')
 141             self.ht = np.trunc(np.float64(l[0])) * np.power(10, np.float64(l[1]))
 142         else:
 143             self.ht = ht
 144         self.K = int(self.time / self.ht)
 145         self.u = np.zeros((3, I + 1, J + 1), 'd')
 146
 147         self._init_bc(0, 0)
 148         self._init_bc(1, 1)
 149         self._init_bc(2, 2)
 150         self._solved = False
 151
 152         if debug:
 153             self._debug_print()
 154
 155     def _debug_print(self):
 156         print "grid: %d x %d" % (self.I, self.J)
 157         print "time:ht:K = %g:%g:%d" % (self.time, self.ht, self.K)
 158         print
 159         print "ly:hy:I = %g:%g:%d" % (self.ly, self.hy, self.I)
 160         print "lz:hz:J = %g:%g:%d" % (self.lz, self.hz, self.J)
 161         print
 162         print "u size (in Kbytes):", ((self.u.size * self.u.itemsize) / 1024.)
 163         print "extrapolated size (in Mbytes):", (float(self.I * self.J * self.K * self.u.itemsize) / (1024**2))
 164         print
 165         print "\gamma_y =", (self.C**2 * self.ht**2) / self.hy**2
 166         print "\gamma_z =", (self.C**2 * self.ht**2) / self.hz**2
 167         print
 168
 169     def _init_bc(self, k, time):
 170         for i in xrange(self.u.shape[1]):
 171             self.u[k, i, 0] = np.sin((np.pi * i) / self.I) * np.sin((2 * np.pi * self.C * self.ht * time) / self.lambda_)
 172
 173     def _make_system(self):
 174         '''Формирует матрицу системы для прогонки (ее диагонали).
 175         '''
 176         J = self.J
 177         d = (self.C**2 * self.ht**2) / self.hz**2
 178         a = np.zeros(J - 2) - d
 179         b = np.zeros(J - 1) + (1 + 2 * d)
 180         c = a.copy()
 181         return a, b, c
 182
 183     def _make_rhs(self, i, k):
 184         '''Формирует правую часть системы для прогонки для шага i по y
 185         '''
 186         g = (self.C**2 * self.ht**2) / self.hy**2
 187         d = (self.C**2 * self.ht**2) / self.hz**2
 188         u = self.u
 189         f = g * (u[1, i + 1, 1:-1] + u[1, i - 1, 1:-1]) + 2 * (1 - g) * u[1, i, 1:-1] - u[0, i, 1:-1]
 190         f[0] += d * u[2, i, 0]
 191         f[-1] += d * u[2, i, self.J]
 192         return f
 193
 194     def _implicit_solve(self):
 195         # система всегда одна и та же
 196         a, b, c = self._make_system()
 197         if not self._solved:
 198             for k in xrange(2, self.K):
 199                 self._init_bc(2, k)
 200                 for i in xrange(1, self.I - 1):
 201                     f = self._make_rhs(i, k)
 202                     self.u[2, i, 1:-1] = native_sweep(a, b, c, f)
 203                 self.u[0] = self.u[1].copy()
 204                 self.u[1] = self.u[2].copy()
 205             self._solver = True
 206         return self.u[2]