Python Function.sub方法代码示例

# 需要导入模块: from dolfin import Function [as 别名]
# 或者: from dolfin.Function import sub [as 别名]

#.........这里部分代码省略.........
      self.Q.apply("add")
    elif self.eigenproblem:
      self.B.apply("add")

                    
  def solve(self, it=0):
    if self.group_GS:
      self.solve_group_GS(it)
    else:
      super(EllipticSNFluxModule, self).solve(it)
      self.slns_mg = split(self.sln)

    i,p,q,k1,k2 = ufl.indices(5)

    sol_timer = Timer("-- Complete solution")
    aux_timer = Timer("---- SOL: Computing angular flux + adjoint")

    # TODO: Move to Discretization
    V11 = FunctionSpace(self.DD.mesh, "CG", self.DD.parameters["p"])

    for gto in range(self.DD.G):
      self.PD.get_xs('D', self.D, gto)

      form = self.D * ufl.diag_vector(ufl.as_matrix(self.DD.ordinates_matrix[i,p]*self.slns_mg[gto][q].dx(i), (p,q)))

      for gfrom in range(self.DD.G):
        pres_Ss = False

        # TODO: Enlarge self.S and self.C to (L+1)^2 (or 1./2.*(L+1)*(L+2) in 2D) to accomodate for anisotropic
        # scattering (lines below using CC, SS are correct only for L = 0, when the following inner loop runs only
        # once.
        for l in range(self.L+1):
          for m in range(-l, l+1):
            if self.DD.angular_quad.get_D() == 2 and divmod(l+m, 2)[1] == 0:
              continue

            pres_Ss |= self.PD.get_xs('Ss', self.S[l], gto, gfrom, l)
            self.PD.get_xs('C', self.C[l], gto, gfrom, l)

        if pres_Ss:
          Cd = ufl.diag(self.C)
          CC = self.tensors.Y[p,k1] * Cd[k1,k2] * self.tensors.Qt[k2,q,i]

          form += ufl.as_vector(CC[p,q,i] * self.slns_mg[gfrom][q].dx(i), p)

      # project(form, self.DD.Vpsi1, function=self.aux_slng, preconditioner_type="petsc_amg")
      # FASTER, but requires form compilation for each dir.:
      for pp in range(self.DD.M):
        assign(self.aux_slng.sub(pp), project(form[pp], V11, preconditioner_type="petsc_amg"))

      self.psi_mg[gto].assign(self.slns_mg[gto] + self.aux_slng)
      self.adj_psi_mg[gto].assign(self.slns_mg[gto] - self.aux_slng)

  def update_phi(self):
    timer = Timer("Scalar flux update")

    for g in range(self.DD.G):
      phig = self.phi_mg[g]
      phig_v = phig.vector()
      psig = self.psi_mg[g]

      for n in range(self.DD.M):
        # This doesn't work (Dolfin Issue #454)
        # assign(self.phi.sub(g), self.phi.sub(g) + self.tensors.Wp[n]*self.psi.sub(g).sub(n))
        psign = psig.sub(n, deepcopy=True)
        phig_v.axpy(float(self.tensors.Wp[n]), psign.vector())

    self.up_to_date["flux"] = True

  def eigenvalue_residual_norm(self,norm_type='l2'):
    if self.group_GS:
      warning("Residual norm can currently be computed only when the whole system is assembled.")
      return 0.
    else:
      super(EllipticSNFluxModule, self).eigenvalue_residual_norm(norm_type)

  def fixed_source_residual_norm(self,norm_type='l2'):
    if self.group_GS:
      warning("Residual norm can currently be computed only when the whole system is assembled.")
      return 0.
    else:
      super(EllipticSNFluxModule,self).fixed_source_residual_norm(norm_type)

  def visualize(self, it=0):
    super(EllipticSNFluxModule, self).visualize()

    labels = ["psi", "adj_psi"]
    functs = [self.psi_mg, self.adj_psi_mg]
    for var,lbl,fnc in zip(["angular_flux", "adjoint_angular_flux"], labels, functs):
      try:
        should_vis = divmod(it, self.parameters["visualization"][var])[1] == 0
      except ZeroDivisionError:
        should_vis = False

      if should_vis:
        for g in range(self.DD.G):
          for n in range(self.DD.M):
            fgn = fnc[g].sub(n)
            fgn.rename(lbl, "{}_g{}_{}".format(lbl, g, n))
            self.vis_files[var][g][n] << (fgn, float(it))

# 需要导入模块: from dolfin import Function [as 别名]
# 或者: from dolfin.Function import sub [as 别名]
    def solve(self, problem):
        self.problem = problem
        doSave = problem.doSave
        save_this_step = False
        onlyVel = problem.saveOnlyVel
        dt = self.metadata['dt']

        nu = Constant(self.problem.nu)
        # TODO check proper use of watches
        self.tc.init_watch('init', 'Initialization', True, count_to_percent=False)
        self.tc.init_watch('rhs', 'Assembled right hand side', True, count_to_percent=True)
        self.tc.init_watch('updateBC', 'Updated velocity BC', True, count_to_percent=True)
        self.tc.init_watch('applybc1', 'Applied velocity BC 1st step', True, count_to_percent=True)
        self.tc.init_watch('applybc3', 'Applied velocity BC 3rd step', True, count_to_percent=True)
        self.tc.init_watch('applybcP', 'Applied pressure BC or othogonalized rhs', True, count_to_percent=True)
        self.tc.init_watch('assembleMatrices', 'Initial matrix assembly', False, count_to_percent=True)
        self.tc.init_watch('solve 1', 'Running solver on 1st step', True, count_to_percent=True)
        self.tc.init_watch('solve 2', 'Running solver on 2nd step', True, count_to_percent=True)
        self.tc.init_watch('solve 3', 'Running solver on 3rd step', True, count_to_percent=True)
        self.tc.init_watch('solve 4', 'Running solver on 4th step', True, count_to_percent=True)
        self.tc.init_watch('assembleA1', 'Assembled A1 matrix (without stabiliz.)', True, count_to_percent=True)
        self.tc.init_watch('assembleA1stab', 'Assembled A1 stabilization', True, count_to_percent=True)
        self.tc.init_watch('next', 'Next step assignments', True, count_to_percent=True)
        self.tc.init_watch('saveVel', 'Saved velocity', True)

        self.tc.start('init')

        # Define function spaces (P2-P1)
        mesh = self.problem.mesh
        self.V = VectorFunctionSpace(mesh, "Lagrange", 2)  # velocity
        self.Q = FunctionSpace(mesh, "Lagrange", 1)  # pressure
        self.PS = FunctionSpace(mesh, "Lagrange", 2)  # partial solution (must be same order as V)
        self.D = FunctionSpace(mesh, "Lagrange", 1)   # velocity divergence space
        if self.bc == 'lagrange':
            L = FunctionSpace(mesh, "R", 0)
            QL = self.Q*L

        problem.initialize(self.V, self.Q, self.PS, self.D)

        # Define trial and test functions
        u = TrialFunction(self.V)
        v = TestFunction(self.V)
        if self.bc == 'lagrange':
            (pQL, rQL) = TrialFunction(QL)
            (qQL, lQL) = TestFunction(QL)
        else:
            p = TrialFunction(self.Q)
            q = TestFunction(self.Q)

        n = FacetNormal(mesh)
        I = Identity(u.geometric_dimension())

        # Initial conditions: u0 velocity at previous time step u1 velocity two time steps back p0 previous pressure
        [u1, u0, p0] = self.problem.get_initial_conditions([{'type': 'v', 'time': -dt},
                                                          {'type': 'v', 'time': 0.0},
                                                          {'type': 'p', 'time': 0.0}])

        if doSave:
            problem.save_vel(False, u0, 0.0)
            problem.save_vel(True, u0, 0.0)

        u_ = Function(self.V)         # current tentative velocity
        u_cor = Function(self.V)         # current corrected velocity
        if self.bc == 'lagrange':
            p_QL = Function(QL)    # current pressure or pressure help function from rotation scheme
            pQ = Function(self.Q)     # auxiliary function for conversion between QL.sub(0) and Q
        else:
            p_ = Function(self.Q)         # current pressure or pressure help function from rotation scheme
        p_mod = Function(self.Q)      # current modified pressure from rotation scheme

        # Define coefficients
        k = Constant(self.metadata['dt'])
        f = Constant((0, 0, 0))

        # Define forms
        # step 1: Tentative velocity, solve to u_
        u_ext = 1.5*u0 - 0.5*u1  # extrapolation for convection term

        # Stabilisation
        h = CellSize(mesh)
        # CBC delta:
        if self.cbcDelta:
            delta = Constant(self.stabCoef)*h/(sqrt(inner(u_ext, u_ext))+h)
        else:
            delta = Constant(self.stabCoef)*h**2/(2*nu*k + k*h*inner(u_ext, u_ext)+h**2)

        if self.use_full_SUPG:
            v1 = v + delta*0.5*k*dot(grad(v), u_ext)
            parameters['form_compiler']['quadrature_degree'] = 6
        else:
            v1 = v

        def nonlinearity(function):
            if self.use_ema:
               return 2*inner(dot(sym(grad(function)), u_ext), v1) * dx + inner(div(function)*u_ext, v1) * dx
                # return 2*inner(dot(sym(grad(function)), u_ext), v) * dx + inner(div(u_ext)*function, v) * dx
                # QQ implement this way?
            else:
                return inner(dot(grad(function), u_ext), v1) * dx

#.........这里部分代码省略.........

# 需要导入模块: from dolfin import Function [as 别名]
# 或者: from dolfin.Function import sub [as 别名]
class TVPD(TV):
    """ Total variation using primal-dual Newton """

    def isPD(self): return True

    def updatePD(self):
        """ Update the parameters.
        parameters should be:
            - k(x) = factor inside TV
            - eps = regularization parameter
            - Vm = FunctionSpace for parameter. 
        ||f||_TV = int k(x) sqrt{|grad f|^2 + eps} dx
        """
        # primal dual variables
        self.Vw = FunctionSpace(self.Vm.mesh(), 'DG', 0)
        self.wH = Function(self.Vw*self.Vw)  # dual variable used in Hessian (re-scaled)
        #self.wH = nabla_grad(self.m)/sqrt(self.fTV) # full Hessian
        self.w = Function(self.Vw*self.Vw)  # dual variable for primal-dual, initialized at 0
        self.dm = Function(self.Vm)
        self.dw = Function(self.Vw*self.Vw)  
        self.testw = TestFunction(self.Vw*self.Vw)
        self.trialw = TrialFunction(self.Vw*self.Vw)
        # investigate convergence of dual variable
        self.dualres = self.w*sqrt(self.fTV) - nabla_grad(self.m)
        self.dualresnorm = inner(self.dualres, self.dualres)*dx
        self.normgraddm = inner(nabla_grad(self.dm), nabla_grad(self.dm))*dx
        # Hessian
        self.wkformPDhess = self.kovsq * ( \
        inner(nabla_grad(self.trial), nabla_grad(self.test)) - \
        0.5*( inner(self.wH, nabla_grad(self.test))*\
        inner(nabla_grad(self.trial), nabla_grad(self.m)) + \
        inner(nabla_grad(self.m), nabla_grad(self.test))*\
        inner(nabla_grad(self.trial), self.wH) ) / sqrt(self.fTV) \
        )*dx
        # update dual variable
        self.Mw = assemble(inner(self.trialw, self.testw)*dx)
        self.rhswwk = inner(-self.w, self.testw)*dx + \
        inner(nabla_grad(self.m)+nabla_grad(self.dm), self.testw) \
        /sqrt(self.fTV)*dx + \
        inner(-inner(nabla_grad(self.m),nabla_grad(self.dm))* \
        self.wH/sqrt(self.fTV), self.testw)*dx

    def compute_dw(self, dm):
        """ Compute dw """
        setfct(self.dm, dm)
        b = assemble(self.rhswwk)
        solve(self.Mw, self.dw.vector(), b)


    def update_w(self, alpha):
        """ update w and re-scale wH """
        self.w.vector().axpy(alpha, self.dw.vector())
        # project each w (coord-wise) onto unit sphere to get wH
        (wx, wy) = self.w.split(deepcopy=True)
        wxa, wya = wx.vector().array(), wy.vector().array()
        normw = np.sqrt(wxa**2 + wya**2)
        factorw = [max(1.0, ii) for ii in normw]
        setfct(wx, wxa/factorw)
        setfct(wy, wya/factorw)
        assign(self.wH.sub(0), wx)
        assign(self.wH.sub(1), wy)
        # check
        (wx,wy) = self.wH.split(deepcopy=True)
        wxa, wya = wx.vector().array(), wy.vector().array()
        assert np.amax(np.sqrt(wxa**2 + wya**2)) <= 1.0 + 1e-14
        # Print results
        dualresnorm = assemble(self.dualresnorm)
        normgraddm = assemble(self.normgraddm)
        print 'line search dual variable: max(|w|)={}, err(w,df)={}, |grad(dm)|={}'.\
        format(np.amax(np.sqrt(normw)), np.sqrt(dualresnorm), np.sqrt(normgraddm))

# 需要导入模块: from dolfin import Function [as 别名]
# 或者: from dolfin.Function import sub [as 别名]
def solve(
        mesh,
        W_element, P_element, Q_element,
        u0, p0, theta0,
        kappa, rho, mu, cp,
        g, extra_force,
        heat_source,
        u_bcs, p_bcs,
        theta_dirichlet_bcs,
        theta_neumann_bcs,
        dx_submesh, ds_submesh
        ):
    # First do a fixed_point iteration. This is usually quite robust and leads
    # to a point from where Newton can converge reliably.
    u0, p0, theta0 = solve_fixed_point(
        mesh,
        W_element, P_element, Q_element,
        theta0,
        kappa, rho, mu, cp,
        g, extra_force,
        heat_source,
        u_bcs, p_bcs,
        theta_dirichlet_bcs,
        theta_neumann_bcs,
        my_dx=dx_submesh,
        my_ds=ds_submesh,
        max_iter=100,
        tol=1.0e-8
        )

    WPQ = FunctionSpace(
        mesh, MixedElement([W_element, P_element, Q_element])
        )
    uptheta0 = Function(WPQ)

    # Initial guess
    assign(uptheta0.sub(0), u0)
    assign(uptheta0.sub(1), p0)
    assign(uptheta0.sub(2), theta0)

    rho_const = rho(_average(theta0))

    stokes_heat_problem = StokesHeat(
        WPQ,
        kappa, rho, rho_const, mu, cp,
        g, extra_force,
        heat_source,
        u_bcs, p_bcs,
        theta_dirichlet_bcs=theta_dirichlet_bcs,
        theta_neumann_bcs=theta_neumann_bcs,
        my_dx=dx_submesh,
        my_ds=ds_submesh
        )

    # solver = FixedPointSolver()
    from dolfin import PETScSNESSolver
    solver = PETScSNESSolver()
    # http://www.mcs.anl.gov/petsc/petsc-current/docs/manualpages/SNES/SNESType.html
    solver.parameters['method'] = 'newtonls'
    # The Jacobian system for Stokes (+heat) are hard to solve.
    # Use LU for now.
    solver.parameters['linear_solver'] = 'lu'
    solver.parameters['maximum_iterations'] = 100
    # TODO tighten tolerance. might not always work though...
    solver.parameters['absolute_tolerance'] = 1.0e-3
    solver.parameters['relative_tolerance'] = 0.0
    solver.parameters['report'] = True

    solver.solve(stokes_heat_problem, uptheta0.vector())

    # u0, p0, theta0 = split(uptheta0)
    # Create a *deep* copy of u0, p0, theta0 to be able to deal with them
    # as actually separate entities vectors.
    u0, p0, theta0 = uptheta0.split(deepcopy=True)
    return u0, p0, theta0