--- a/src/convection_diffusion.py Mon Jun 29 00:20:42 2026 +0300
+++ b/src/convection_diffusion.py Sun Oct 04 13:02:26 2026 -0500
@@ -64,7 +64,7 @@
u0,
g,
# w0,
- domain_size=[-2, 2, -2, 2],
+ domain_dims=[[-2, 2], [-2, 2]],
nx=32,
ny=32,
degree=2,
@@ -89,11 +89,11 @@
self.C_emb = C_emb # embedding constant
self.C_reg = C_reg
if Delta_t is None:
- self.Delta_t = T / num_steps # Matches your N=50 intent
+ self.Delta_t = T / num_steps
else:
self.Delta_t = Delta_t
- self.domain_size = domain_size
+ self.domain_dims = domain_dims
self.nx = nx
self.ny = ny
self.degree = degree
@@ -112,19 +112,18 @@
domain = mesh.create_rectangle(
MPI.COMM_SELF,
[
- np.array([domain_size[0], domain_size[2]]),
- np.array([domain_size[1], domain_size[3]]),
+ np.array([domain_dims[0][0], domain_dims[1][0]]),
+ np.array([domain_dims[0][1], domain_dims[1][1]]),
],
[nx, ny],
mesh.CellType.triangle,
)
+ # Compute the domain volume for the adjoint Lipschitz factor
self.domain = domain
- # Compute the domain volume for the adjoint Lipschitz factor
- self.domain_size = domain_size
- # Domain diameter from your domain_size
+ # Domain diameter from your domain
self.diam = np.sqrt(
- (domain_size[1] - domain_size[0]) ** 2
- + (domain_size[3] - domain_size[2]) ** 2
+ (domain_dims[0][0] - domain_dims[0][1]) ** 2
+ + (domain_dims[1][0] - domain_dims[1][1]) ** 2
)
V = fem.functionspace(domain, ("Lagrange", degree))
@@ -140,26 +139,36 @@
self.bcs = [] # Initialize BCS as LIST
self.adjoint_bcs = [] # Initialize ADJOINT BCS as LIST (w = 0)
- # Identify boundary facets for Dirichlet BC location (e.g. left/right boundaries at x=domain_size[0] or x=domain_size[1])
+ # Identify boundary facets for Dirichlet BC location (e.g. left/right boundaries
fdim = domain.topology.dim - 1
# Γ₁ facets: TOP/BOTTOM boundaries (y = ymin, ymax) - STATE Dirichlet u = g
gamma1_facets = mesh.locate_entities_boundary(
domain,
fdim,
lambda x: (
- np.isclose(x[1], domain_size[2]) # x[1] = y = ymin
- | np.isclose(x[1], domain_size[3])
+ np.isclose(x[1], domain_dims[1][0]) # x[1] = y = ymin
+ | np.isclose(x[1], domain_dims[1][1])
), # x[1] = y = ymax
)
- # Γ₂ facets: LEFT/RIGHT boundaries (x = xmin, xmax) - ADJOINT Dirichlet w = 0
+ # Γ₂ facets: LEFT/RIGHT boundaries (x = xmin, xmax)
gamma2_facets = mesh.locate_entities_boundary(
domain,
fdim,
lambda x: (
- np.isclose(x[0], domain_size[0]) # x[0] = x = xmin
- | np.isclose(x[0], domain_size[1])
+ np.isclose(x[0], domain_dims[0][0]) # x[0] = x = xmin
+ | np.isclose(x[0], domain_dims[0][1])
), # x[0] = x = xmax
)
+ # tag Γ2 facets (left/right) once in __init__
+ facet_idx = np.concatenate([gamma1_facets, gamma2_facets])
+ facet_vals = np.concatenate(
+ [
+ np.full(len(gamma1_facets), 1, np.int32),
+ np.full(len(gamma2_facets), 2, np.int32),
+ ]
+ )
+ order = np.argsort(facet_idx)
+ facet_tags = mesh.meshtags(domain, fdim, facet_idx[order], facet_vals[order])
# STATE BC: u = g on Γ₁ (top/bottom)
dofs_gamma1 = fem.locate_dofs_topological(V, fdim, gamma1_facets)
# Create BC and ADD TO LIST
@@ -168,12 +177,12 @@
# ADJOINT BC: w = 0 on Γ₁ (same facets, homogenized)
zero = fem.Constant(domain, PETSc.ScalarType(0.0))
- dofs_gamma2 = fem.locate_dofs_topological(V, fdim, gamma2_facets)
- bc_adjoint = fem.dirichletbc(zero, dofs_gamma2, V)
+ # dofs_gamma2 = fem.locate_dofs_topological(V, fdim, gamma2_facets)
+ bc_adjoint = fem.dirichletbc(zero, dofs_gamma1, V)
self.adjoint_bcs.append(bc_adjoint)
self.fdim = fdim
- self.boundary_facets = np.unique(np.concatenate((gamma1_facets, gamma2_facets)))
+ # self.boundary_facets = np.unique(np.concatenate((gamma1_facets, gamma2_facets)))
ux = fem.Function(V)
wpx = fem.Function(V)
@@ -211,21 +220,21 @@
j_u = fem.Function(V) # New copy!
j_u.name = "j_u"
self.j_u = j_u
+ ds = ufl.Measure("ds", domain=domain, subdomain_data=facet_tags)
+ n = ufl.FacetNormal(domain)
a2 = (
w * v * ufl.dx
- + dt
- * (
- ufl.dot(self.k * ufl.grad(w), ufl.grad(v))
- - (self.c0 * w.dx(0) * v + self.c1 * w.dx(1) * v)
- )
- * ufl.dx
+ + dt * ufl.dot(self.k * ufl.grad(w), ufl.grad(v)) * ufl.dx
+ - dt * (self.c0 * w.dx(0) + self.c1 * w.dx(1)) * v * ufl.dx
+ + dt * (self.c0 * n[0] + self.c1 * n[1]) * w * v * ds(2)
)
self.bilinear_form_w = fem.form(a2)
# Step n=N-1: L2 uses w_n = w^N=0 → solves w^{N-1}. (∂t)_primal^*
# Step n=N-2: L2 uses w_n = w^{N-1} → solves w^{N-2}
# Replace your L2 with 2 simple forms (same notation)
- L2 = w_n * v * ufl.dx + dt * j_u * v * ufl.dx # kown w_n (previous step)
+ # No dt scaling in j_u, as that's already included in it.
+ L2 = w_n * v * ufl.dx + j_u * v * ufl.dx # kown w_n (previous step)
self.linear_form_w = fem.form(L2)
# Solving forward PDE
@@ -290,7 +299,7 @@
b_ps.setValuesLocal(
cell_dofs,
- alpha * phi,
+ alpha * phi * dt,
addv=PETSc.InsertMode.ADD_VALUES,
)
@@ -351,8 +360,6 @@
assert len(j) == num_steps, f"MISMATCH: j({len(j)}) != u_n_list({num_steps})"
# Terminal condition: w(T) = 0
w_n = self.w_n
- w_n.x.array[:] = 0.0
- w_n.x.scatter_forward()
self.c0.value = c[0]
self.c1.value = c[1]
@@ -375,11 +382,12 @@
t = self.t0 + num_steps * dt # Start at time T
# N+1 adjoint values to match N PDE steps
- w_n_list = [None] * (num_steps + 1)
+ w_n_list = [None] * num_steps
whT = fem.Function(V)
whT.x.array[:] = 0.0 # w^N = 0
whT.x.scatter_forward()
- w_n_list[num_steps] = whT # Store at FINAL index
+ w_n.x.array[:] = whT.x.array
+ w_n.x.scatter_forward()
# backward steps**: compute w^{N-1}, ..., w^0
@@ -436,15 +444,7 @@
# 1. expr1: dual sensitivity w.r.t. μ = ∫₀^T w_p dt (always Function)
wp_bar = fem.Function(V)
- wp_bar.x.array[:] = (
- 0.5
- * dt
- * (
- w_n_list[0].x.array
- + w_n_list[-1].x.array
- + 2 * np.sum([w.x.array for w in w_n_list[1:-1]], axis=0)
- )
- )
+ wp_bar.x.array[:] = dt * sum(w.x.array for w in w_n_list)
ux = self.ux
wpx = self.wpx
@@ -462,8 +462,7 @@
wpx.x.array[:] = wp.x.array[:]
wpx.x.scatter_forward()
v2 = fem.assemble_scalar(form)
- weight = 0.5 * dt if (n == 0 or n == len(u_n_list) - 1) else dt
- expr2 -= weight * v2
+ expr2 -= dt * v2
else:
raise Exception("Unimplemented - out of date")
@@ -475,8 +474,7 @@
for n, (u, wp) in enumerate(zip(u_n_list, w_n_list)):
form = ufl.inner(ufl.grad(u), ufl.grad(wp)) * ufl.dx
v2 = fem.assemble_scalar(fem.form(form))
- weight = 0.5 * dt if (n == 0 or n == len(u_n_list) - 1) else dt
- expr2a -= weight * v2
+ expr2a -= dt * v2
# 3. expr3: dual sensitivity w.r.t. c = (c1, c2)
if is_scalar_c:
@@ -492,8 +490,7 @@
wpx.x.array[:] = wp.x.array[:]
wpx.x.scatter_forward()
v3 = fem.assemble_scalar(forms[i])
- weight = 0.5 * dt if (n == 0 or n == len(u_n_list) - 1) else dt
- expr3i -= weight * v3
+ expr3i -= dt * v3
return expr3i
expr3 = tuple(val(i) for i in range(dim))
@@ -509,8 +506,7 @@
for n, (u, wp) in enumerate(zip(u_n_list, w_n_list)):
form = wp * ufl.grad(u)[i] * ufl.dx
v3 = fem.assemble_scalar(fem.form(form))
- weight = 0.5 * dt if (n == 0 or n == len(u_n_list) - 1) else dt
- expr3ia -= weight * v3
+ expr3ia -= dt * v3
return expr3i
expr3 = tuple(val(i) for i in range(dim))
@@ -781,8 +777,9 @@
def _own(self, x):
return x if isinstance(x, float) else own(x)
- def __init__(self, α, base=None):
+ def __init__(self, α, base=None, τscale=None):
self.α = α
+ self.τscale = τscale
if base is not None:
(k, (c1, c2)) = base
self.base = (self._own(k), (self._own(c1), self._own(c2)))
@@ -846,6 +843,11 @@
else:
αk, αc = self.α
+ if self.τscale is not None:
+ kscale, cscale = self.τscale
+ αk = αk * kscale
+ αc = αc * cscale
+
return (
self._prox1(τ * αk, k, kb),
(self._prox1(τ * αc, c1, c1b), self._prox1(τ * αc, c2, c2b)),
@@ -995,9 +997,9 @@
class BoxedQuadraticRegularisation:
- def __init__(self, ω0, ω1, α, base=None):
+ def __init__(self, ω0, ω1, α, base=None, τscale=None):
self.box = BoxConstraints(ω0, ω1)
- self.quadratic = QuadraticRegularisation(α, base)
+ self.quadratic = QuadraticRegularisation(α, base, τscale)
def apply(self, x):
return self.box.apply(x) + self.quadratic.apply(x)