diff -r c4f923017349 -r a0d61ac3719e src/convection_diffusion.py --- 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)