sweep.py

   1 #!/usr/bin/python
   2 #-*- coding: utf-8 -*-
   3
   4 import numpy as np
   5 from scipy.weave import blitz, inline, converters
   6
   7
   8 def sweep(_a, _b, _c, _d):
   9     n = len(_b)
  10     a = np.zeros(n)
  11     a[1:] = _a
  12     b = _b.copy()
  13     c = np.zeros(n)
  14     c[:-1] = _c
  15     x = _d.copy()
  16
  17     # Прямой ход
  18     c[0] /= b[0]
  19     x[0] /= b[0]
  20     for i in xrange(1, n):
  21         denom = 1 / (b[i] - c[i-1] * a[i])
  22         c[i] *= denom
  23         x[i] = (x[i] - x[i-1] * a[i]) * denom
  24
  25     # Обратный ход
  26     for i in xrange(n - 2, -1, -1):
  27         x[i] -= c[i] * x[i+1]
  28
  29     return x
  30
  31 def native_sweep(_a, _b, _c, _d):
  32     n = len(_b)
  33     a = np.zeros(n)
  34     a[1:] = _a
  35     b = _b.copy()
  36     c = np.zeros(n)
  37     c[:-1] = _c
  38     x = _d.copy()
  39
  40     code = """
  41     c(0) /= b(0);
  42     x(0) /= b(0);
  43     for (int i=1; i < n; i++) {
  44         double denom = 1 / (b(i) - c(i - 1) * a(i));
  45         c(i) *= denom;
  46         x(i) = (x(i) - x(i - 1) * a(i)) * denom;
  47     }
  48     for (int i=n-2; i >= 0; i--) {
  49         x(i) -= c(i) * x(i+1);
  50     }
  51     """
  52     inline(code,
  53            ['n', 'a', 'b', 'c', 'x'],
  54            type_converters=converters.blitz,
  55            compiler='gcc')
  56     return x
  57
  58 def main():
  59     pass
  60
  61
  62 if __name__ == '__main__':
  63     main()