SHARE
TWEET

Untitled

a guest Jun 20th, 2019 47 Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. '''
  2. Simple example coupling a serial and distributed ImplicitComponent
  3. '''
  4.  
  5. import numpy as np
  6.  
  7. import openmdao.api as om
  8. from mpi4py import MPI
  9. from openmdao.utils.array_utils import evenly_distrib_idxs
  10.  
  11. rank = MPI.COMM_WORLD.rank
  12.  
  13. size = 3
  14. A = np.array([[1.0, 8.0, 0.0], [-1.0, 10.0, 2.0], [3.0, 100.5, 1.0]])
  15.  
  16. '''
  17. This component solves the following quadratic equation in parallel:
  18.     a_i0 * y_i^2 + a_i1 * y_i + a_i2 = x_i
  19.     for i = {0,1,2}
  20. where the coefficients are the components of the matrix A
  21. '''
  22. class DistribQuadtric(om.ImplicitComponent):
  23.     def initialize(self):
  24.         self.options['distributed'] = True
  25.  
  26.         self.options.declare('size', types=int, default=1,
  27.                              desc="Size of input and output vectors.")
  28.  
  29.     def setup(self):
  30.         comm = self.comm
  31.         rank = comm.rank
  32.  
  33.         size_total = self.options['size']
  34.  
  35.         # Distribute x and y vectors across each processor as evenly as possible
  36.         sizes, offsets = evenly_distrib_idxs(comm.size, size_total)
  37.         start = offsets[rank]
  38.         end = start + sizes[rank]
  39.     self.size_local = size_local = sizes[rank]
  40.  
  41.     # Get the local slice of A that this processor will be working with
  42.     self.A_local = A[start:end,:]
  43.  
  44.         self.add_input('x', np.ones(size_local, float),
  45.                        src_indices=np.arange(start, end, dtype=int))
  46.  
  47.         self.add_output('y', np.ones(size_local, float))
  48.  
  49.     def apply_nonlinear(self, inputs, outputs, residuals):
  50.     x = inputs['x']
  51.     y = outputs['y']
  52.     r = residuals['y']
  53.         for i in range(self.size_local):
  54.         r[i] = self.A_local[i, 0] * y[i]**2 + self.A_local[i, 1] * y[i]
  55.         + self.A_local[i, 2] - x[i]
  56.  
  57.     def solve_nonlinear(self, inputs, outputs):
  58.     x = inputs['x']
  59.     y = outputs['y']
  60.         for i in range(self.size_local):
  61.         a = self.A_local[i, 0]
  62.         b = self.A_local[i, 1]
  63.         c = self.A_local[i, 2] - x[i]
  64.         y[i] = (-b + np.sqrt(b**2 - 4*a*c))/(2*a)
  65.  
  66. '''
  67. This component solves the following linear equation in serial:
  68.     Ax = y
  69. '''
  70. class SerialLinear(om.ImplicitComponent):
  71.     def initialize(self):
  72.  
  73.         self.options.declare('size', types=int, default=1,
  74.                              desc="Size of input and output vectors.")
  75.  
  76.     def setup(self):
  77.     size = self.options['size']
  78.  
  79.         self.add_input('y', np.ones(size, float))
  80.  
  81.         self.add_output('x', np.ones(size, float))
  82.  
  83.     self.A = A
  84.  
  85.     def apply_nonlinear(self, inputs, outputs, residuals):
  86.     y = inputs['y']
  87.     x = outputs['x']
  88.     r = residuals['x']
  89.         r = y - A.dot(x)
  90.  
  91.     def solve_nonlinear(self, inputs, outputs):
  92.     y = inputs['y']
  93.     x = outputs['x']
  94.         x[:] = np.linalg.inv(A).dot(y)
  95.  
  96. # Create a couple problem between the linear and quadratic components
  97. prob = om.Problem()
  98. top_group = prob.model
  99. top_group.add_subsystem("distributed_quad", DistribQuadtric(size=size))
  100. top_group.add_subsystem("serial_linear", SerialLinear(size=size))
  101.  
  102. # Connect variables between components
  103. top_group.connect('serial_linear.x', 'distributed_quad.x')
  104. top_group.connect('distributed_quad.y', 'serial_linear.y')
  105.  
  106. # Need a nonlinear solver since the model is coupled
  107. top_group.nonlinear_solver = om.NonlinearBlockGS(iprint=2, maxiter=20)
  108.  
  109. # Setup problem
  110. prob.setup()
  111.  
  112. # Solver problem
  113. prob.run_model()
  114.  
  115. # Print out solution
  116. if prob.comm.rank == 0:
  117.     print('x', prob['serial_linear.x'])
  118.     print('y', prob['serial_linear.y'])
  119.      
  120. NL: NLBGS 0 ; 2.35754338 1
  121. NL: NLBGS 1 ; 0.256315721 0.108721529
  122. NL: NLBGS 2 ; 0.036527896 0.0154940504
  123. NL: NLBGS 3 ; 0.00641965062 0.00272302545
  124. NL: NLBGS 4 ; 0.0011292331 0.000478987198
  125. NL: NLBGS 5 ; 0.000198654857 8.42635002e-05
  126. NL: NLBGS 6 ; 3.49479079e-05 1.48238663e-05
  127. NL: NLBGS 7 ; 6.14814792e-06 2.60786205e-06
  128. NL: NLBGS 8 ; 1.08160237e-06 4.58783657e-07
  129. NL: NLBGS 9 ; 1.90279057e-07 8.0710734e-08
  130. NL: NLBGS 10 ; 3.34745201e-08 1.41988989e-08
  131. NL: NLBGS 11 ; 5.8889481e-09 2.49791717e-09
  132. NL: NLBGS 12 ; 1.03600386e-09 4.3944212e-10
  133. NL: NLBGS 13 ; 1.8225669e-10 7.7307884e-11
  134. NL: NLBGS Converged
  135. ('x', array([-0.01251987,  0.00136932, -0.11111688]))
  136. ('y', array([-0.00156529, -0.19602066, -0.01105954]))
  137.      
  138. NL: NLBGS 0 ; 5.66931072 1
  139. NL: NLBGS 1 ; 0.6855401 0.120921243
  140. NL: NLBGS 2 ; 0.0993351375 0.0175215546
  141. NL: NLBGS 3 ; 0.0174731006 0.00308205026
  142. NL: NLBGS 4 ; 0.00307353315 0.000542135243
  143. NL: NLBGS 5 ; 0.00054069662 9.537255e-05
  144. NL: NLBGS 6 ; 9.51208366e-05 1.67782013e-05
  145. NL: NLBGS 7 ; 1.67339624e-05 2.95167495e-06
  146. NL: NLBGS 8 ; 2.94389363e-06 5.19268351e-07
  147. NL: NLBGS 9 ; 5.17899477e-07 9.1351401e-08
  148. NL: NLBGS 10 ; 9.11105862e-08 1.60708401e-08
  149. NL: NLBGS 11 ; 1.60284752e-08 2.82723526e-09
  150. NL: NLBGS 12 ; 2.81978416e-09 4.97376895e-10
  151. NL: NLBGS 13 ; 4.96064272e-10 8.74999266e-11
  152. NL: NLBGS Converged
  153. ('x', array([-0.01251987,  0.00136932, -0.11111688]))
  154. ('y', array([-0.00156529, -0.19602066, -0.01105954]))
  155.      
  156. def get_norm(self):
  157.         """
  158.         Return the norm of this vector.
  159.  
  160.         Returns
  161.         -------
  162.         float
  163.             norm of this vector.
  164.         """
  165.         return self._system.comm.allreduce(np.linalg.norm(self._data))
RAW Paste Data
We use cookies for various purposes including analytics. By continuing to use Pastebin, you agree to our use of cookies as described in the Cookies Policy. OK, I Understand
Not a member of Pastebin yet?
Sign Up, it unlocks many cool features!
 
Top