heat_complete.py

来自「利用有限元算法计算各类微分方程」· Python 代码 · 共 104 行

from dolfin import *

# Source term
class Source(Function):
    def __init__(self, element, mesh):
        Function.__init__(self, element, mesh)
    def eval(self, values, x):
	values[0] = 0.0

# Source term
class InitialValue(Function):
    def __init__(self, element, mesh):
        Function.__init__(self, element, mesh)
    def eval(self, values, x):
	values[0] = 1.0

# Neumann boundary condition
class Flux(Function):
    def __init__(self, element, mesh):
        Function.__init__(self, element, mesh)
    def eval(self, values, x):
        values[0] = 0.0

# Sub domain for Dirichlet boundary condition
class DirichletBoundary(SubDomain):
    def inside(self, x, on_boundary):
        return bool(on_boundary)

# Load mesh and create finite elements
mesh = Mesh("Square.xml")
mesh.refine()
element = FiniteElement("Lagrange", "triangle", 1)
element2 = FiniteElement("Lagrange", "triangle", 2)

# Define variational problem
v = TestFunction(element)
U = TrialFunction(element)
f = Source(element2, mesh)
g = Flux(element, mesh)
c = 0.01

# Sub domain for Dirichlet boundary condition
class DirichletBoundary(SubDomain):
    def inside(self, x, on_boundary):
        return bool(on_boundary)

# Initialise source function and previous solution function
f = Source(element2, mesh)
u0 = InitialValue(element, mesh)

# Parameters
T = 0.5
k = 0.01
t = k

# Test and trial functions
v = TestFunction(element)
u = TrialFunction(element)

# Functions
x0 = Vector()
x1 = Vector()
u1 = Function(element, mesh, x1)
h = MeshSize("triangle", mesh)
c = 1.0


# Variational problem
a = (u*v+k*dot(grad(u),grad(v)))*dx
L = (u0*v+f*v)*dx

# Set up boundary condition
boundary = DirichletBoundary()
zero = Function(mesh, 0.0)
bc = DirichletBC(zero, mesh, boundary)

# Assemble matrix (since heat coefficient c is constant in time, the matrix
# A is also constant in time)
A = assemble(a, mesh)


# Output file
out_file = File("temperature.pvd")

# Time-stepping
while t < T:

    # Assemble vector and apply boundary conditions
    b = assemble(L, mesh)
    bc.apply(A, b, a)
    
    # Solve the linear system
    solve(A, x1, b)
    
    # Save the solution to file
    out_file << u1

    # Move to next interval
    t += k

    # Copy solution from previous interval
    u0.assign(u1)