Mercurial > repos > pointsource_pde / file comparison

comparison: src/convection_diffusion.py

src/convection_diffusion.py

changeset 3
c3a4f4bb87f7
parent 1
a4137aedcb3a
equal deleted inserted replaced
1:a4137aedcb3a 3:c3a4f4bb87f7
76 C_stab=1.0, 76 C_stab=1.0,
77 alpha=1.0, # \alpha = k_m/2, k convective coefficient, k>=km 77 alpha=1.0, # \alpha = k_m/2, k convective coefficient, k>=km
78 beta=1.0, # \beta = \| c \|_\infty^2 /k_m 78 beta=1.0, # \beta = \| c \|_\infty^2 /k_m
79 C_emb=1.0, 79 C_emb=1.0,
80 C_reg=1.0, 80 C_reg=1.0,
81 xbound=XBound,
82 override_lipschitz=None,
83 override_lipschitz_pair=None,
81 ): 84 ):
82 self.p = p # L^p, p=2 85 self.p = p # L^p, p=2
83 self.C_stab = C_stab # adjoint stability constant 86 self.C_stab = C_stab # adjoint stability constant
84 self.alpha = alpha # coercivity lower bound for k 87 self.alpha = alpha # coercivity lower bound for k
85 self.beta = beta # theoretical parameter 88 self.beta = beta # theoretical parameter
93 self.domain_size = domain_size 96 self.domain_size = domain_size
94 self.nx = nx 97 self.nx = nx
95 self.ny = ny 98 self.ny = ny
96 self.degree = degree 99 self.degree = degree
97 self.T = T 100 self.T = T
101 self.xbound = xbound
98 102
99 self.num_steps = num_steps 103 self.num_steps = num_steps
100 self.dt = (T - t0) / num_steps 104 self.dt = (T - t0) / num_steps
101 self.t0 = t0 105 self.t0 = t0
102 106
103 self.save_for_plot = False 107 self.save_for_plot = False
108
109 self.override_lipschitz = override_lipschitz
110 self.override_lipschitz_pair = override_lipschitz_pair
104 111
105 domain = mesh.create_rectangle( 112 domain = mesh.create_rectangle(
106 MPI.COMM_SELF, 113 MPI.COMM_SELF,
107 [ 114 [
108 np.array([domain_size[0], domain_size[2]]), 115 np.array([domain_size[0], domain_size[2]]),
245 bilinear_form = self.bilinear_form 252 bilinear_form = self.bilinear_form
246 A = assemble_matrix(bilinear_form, bcs=bcs) 253 A = assemble_matrix(bilinear_form, bcs=bcs)
247 A.assemble() 254 A.assemble()
248 255
249 # Create reusable RHS vector (will be filled each timestep) 256 # Create reusable RHS vector (will be filled each timestep)
250 b = create_vector(linear_form) 257 b = create_vector(linear_form.function_spaces)
251 b_ps = create_vector(linear_form) 258 b_ps = create_vector(linear_form.function_spaces)
252 259
253 # Prepare solver once 260 # Prepare solver once
254 solver = PETSc.KSP().create(domain.comm) 261 solver = PETSc.KSP().create(domain.comm)
255 solver.setOperators(A) 262 solver.setOperators(A)
256 solver.setType(PETSc.KSP.Type.PREONLY) 263 solver.setType(PETSc.KSP.Type.PREONLY)
356 363
357 A2 = assemble_matrix(bilinear_form_w, bcs=adjoint_bcs) # Fixed bcs 364 A2 = assemble_matrix(bilinear_form_w, bcs=adjoint_bcs) # Fixed bcs
358 A2.assemble() 365 A2.assemble()
359 366
360 # Create vector for RHS to be updated each step 367 # Create vector for RHS to be updated each step
361 b2 = create_vector(linear_form_w) 368 b2 = create_vector(linear_form_w.function_spaces)
362 369
363 solver = PETSc.KSP().create(domain.comm) 370 solver = PETSc.KSP().create(domain.comm)
364 solver.setOperators(A2) 371 solver.setOperators(A2)
365 solver.setType(PETSc.KSP.Type.PREONLY) 372 solver.setType(PETSc.KSP.Type.PREONLY)
366 solver.getPC().setType(PETSc.PC.Type.LU) 373 solver.getPC().setType(PETSc.PC.Type.LU)
511 if self.save_for_plot: 518 if self.save_for_plot:
512 self.plot_wp_bar = wp_bar 519 self.plot_wp_bar = wp_bar
513 self.plot_u_n_list = u_n_list 520 self.plot_u_n_list = u_n_list
514 self.plot_w_n_list = w_n_list 521 self.plot_w_n_list = w_n_list
515 self.plot_j_n_list = j 522 self.plot_j_n_list = j
523 self.plot_aux = (k, c)
516 524
517 return wp_bar, (expr2, expr3) 525 return wp_bar, (expr2, expr3)
518 526
519 def do_plot(self, iter, μ, frames, prefix): 527 def do_plot(self, iter, μ, frames, prefix):
520 n = self.num_steps 528 n = self.num_steps
529 k, (c1, c2) = self.plot_aux
521 np.savez_compressed( 530 np.savez_compressed(
522 "%s/res_%d.npz" % (prefix, iter), 531 "%s/res_%d.npz" % (prefix, iter),
523 wp_bar_array=self.plot_wp_bar.x.array, 532 wp_bar_array=self.plot_wp_bar.x.array,
524 u_n_list_array=[self.plot_u_n_list[i].x.array for i in frames], 533 u_n_list_array=[self.plot_u_n_list[i].x.array for i in frames],
525 w_n_list_array=[self.plot_w_n_list[i].x.array for i in frames], 534 w_n_list_array=[self.plot_w_n_list[i].x.array for i in frames],
526 j_n_list_array=[self.plot_j_n_list[i].x.array for i in frames], 535 j_n_list_array=[self.plot_j_n_list[i].x.array for i in frames],
536 k=k,
537 c1=c1,
538 c2=c2,
527 frames=frames, 539 frames=frames,
528 times=list(map(lambda i: self.T * i / (n - 1), frames)), 540 times=list(map(lambda i: self.T * i / (n - 1), frames)),
529 mu=[[point[0], point[1], alpha] for point, alpha in μ.iter_padded()], 541 mu=[[point[0], point[1], alpha] for point, alpha in μ.iter_padded()],
530 ) 542 )
531 543 # print("Aux variables: k: %f, c1: %f, c2: %f" % (k, c1, c2))
532 def _stability_prefactor( 544
533 self, *, C_stab=None, alpha=None, beta=None, T=None, C_emb=None, C_reg=None 545 # def _stability_prefactor(
534 ): 546 # self, *, C_stab=None, alpha=None, beta=None, T=None, C_emb=None, C_reg=None
535 """ 547 # ):
536 Return the stability prefactor 548 # """
537 P := (C_stab/alpha) * (1 + sqrt(T)*(1 + sqrt(beta)/sqrt(alpha))) 549 # Return the stability prefactor
538 and the measure-to-functional constant L_e_mu := C_emb * C_reg * sqrt(T). 550 # P := (C_stab/alpha) * (1 + sqrt(T)*(1 + sqrt(beta)/sqrt(alpha)))
539 Return (P, L_e_mu): 551 # and the measure-to-functional constant L_e_mu := C_emb * C_reg * sqrt(T).
540 """ 552 # Return (P, L_e_mu):
541 # self._stability_prefactor(alpha=0.5, T=2.0) 553 # """
542 554 # # self._stability_prefactor(alpha=0.5, T=2.0)
543 # Read from arguments or instance attributes 555
544 C_stab = C_stab if C_stab is not None else self.C_stab 556 # # Read from arguments or instance attributes
545 alpha = alpha if alpha is not None else self.alpha 557 # C_stab = C_stab if C_stab is not None else self.C_stab
546 beta = beta if beta is not None else self.beta 558 # alpha = alpha if alpha is not None else self.alpha
547 T = T if T is not None else self.T 559 # beta = beta if beta is not None else self.beta
548 C_emb = C_emb if C_emb is not None else self.C_emb 560 # T = T if T is not None else self.T
549 C_reg = C_reg if C_reg is not None else self.C_reg 561 # C_emb = C_emb if C_emb is not None else self.C_emb
550 562 # C_reg = C_reg if C_reg is not None else self.C_reg
551 # Stability factor P 563
552 564 # # Stability factor P
553 P = (C_stab / alpha) * (1 + np.sqrt(T) * (1 + np.sqrt(beta) / np.sqrt(alpha))) 565
554 L_e_mu = C_emb * C_reg * np.sqrt(T) 566 # P = (C_stab / alpha) * (1 + np.sqrt(T) * (1 + np.sqrt(beta) / np.sqrt(alpha)))
555 567 # L_e_mu = C_emb * C_reg * np.sqrt(T)
556 return P, L_e_mu 568
557 569 # return P, L_e_mu
558 # Solution operator Lipschitz factor 570
559 def lipschitz_factor(self): 571 # # Solution operator Lipschitz factor
560 """ 572 # def lipschitz_factor(self, xbound=None):
561 Conservative Lipschitz factor for the forward solution operator mapping 573 # """
562 (k,c,mu) -> u. Returns a single scalar C_Lip such that 574 # Conservative Lipschitz factor for the forward solution operator mapping
563 ||u2 - u1|| <= C_Lip * (||k2-k1|| + ||c2-c1|| + ||mu2-mu1||_* ) 575 # (k,c,mu) -> u. Returns a single scalar C_Lip such that
564 (we use a conservative form combining the measure and (k,c) pieces). 576 # ||u2 - u1|| <= C_Lip * (||k2-k1|| + ||c2-c1|| + ||mu2-mu1||_* )
565 """ 577 # (we use a conservative form combining the measure and (k,c) pieces).
566 P, L_e_mu = self._stability_prefactor() 578 # """
567 # Conservative combined choice: multiply P by max(1, L_e_mu) 579 # P, L_e_mu = self._stability_prefactor()
568 C_Lip = P * max(1.0, L_e_mu) 580 # # Conservative combined choice: multiply P by max(1, L_e_mu)
569 return C_Lip 581 # C_Lip = P * max(1.0, L_e_mu)
570 582 # return C_Lip
571 # Adjoint solution operator Lipschitz factor 583 def diff_chain_lipschitz_factor(self, Delta_t=None):
572 """ 584 """
573 def diff_adj_lipschitz_factor(self): 585 Compute a Lipschitz factor C_tot^Lip using locally accumulated bounds
574 586 M_J = 1.
575 Compute adjoint-solution Lipschitz factor A for 2D domain. 587 """
576 588
577 For p = 2, A = 1. 589 if self.override_lipschitz is not None:
578 For p > 2, A = |Omega|^(1/2 - 1/p), where |Omega| is the mesh area. 590 return self.override_lipschitz
579 591
580 p = self.p 592 # Use existing diff_bound3
581 omega_vol = self.domain_volume 593 C_mu_adj, C_k_adj, C_c_adj = self.diff_bound3(Delta_t=Delta_t)
582 return 1.0 if p == 2 else omega_vol ** (0.5 - 1.0 / p) 594
583 """ 595 if Delta_t is None:
584 596 # fall back to the scheme used in diff_bound3
585 def diff_adj_lipschitz_factor(self): 597 Delta_t = self.dt
586 """ 598
587 Always return 1 for p=2 599 N = int(np.ceil(self.T / Delta_t))
588 """ 600
589 return 1.0 601 # Use local per‑step bounds
590 602 C_adj_step = C_mu_adj / N
591 def diff_bound(self, xbound=None): 603 C_state_step = C_k_adj / N
592 """ 604 C_c_step = C_c_adj / N
593 Differential bound for the solution operator derivative. 605
594 606 # Per‑step "Lipschitz"
595 Returns the tuple: (L_e_mu, L_e_k, L_e_c) 607 l_e_local = C_adj_step + C_state_step + C_c_step
596 608
597 where 609 # Global Lipschitz constants
598 L_e_k = ||∇u||_{L^2(Ω_T)} 610 L_e = l_e_local * N
599 L_e_c = ||∇u||_{L^2(Ω_T)} 611 M_e = 1.2 * L_e
600 L_e_mu = C_emb * C_reg * sqrt(T) 612 L_A = 0.5 * C_state_step * N
601 """ 613
602 614 # PDE/coercivity
603 # If grad_u_norm is None, compute it from u0 615 alpha = max(0.1, self.alpha) # guard α ≥ 0.1
604 grad_u_norm = getattr(self, "grad_u_norm", None) 616 M_J = 1.0
605 if grad_u_norm is None: 617
606 u = self.u0 618 # C_tot^Lip
607 V = self.V 619 inv_alpha = 1.0 / alpha
608 # Compute L2 norm of gradient of u0 over the domain 620 inv_alpha_sq = inv_alpha * inv_alpha
609 grad_u = ufl.grad(u) 621 C_tot = inv_alpha * (L_e + inv_alpha_sq * L_A * M_e) * M_J
610 grad_u_sq = ufl.inner(grad_u, grad_u) 622 return C_tot
611 dx = ufl.Measure("dx", domain=self.domain) 623
612 # Integral of |∇u|^2 over Ω 624 def diff_chain_lipschitz_factor_pair(self, Delta_t=None):
613 grad_u_norm = fem.assemble_scalar(fem.form(grad_u_sq * dx)) ** 0.5 625 """
614 # Optionally store it for later reuse 626 Return (C^Lip_μ/M_J, C^Lip_aux/M_J)
615 self.grad_u_norm = grad_u_norm 627 """
616 628
617 return grad_u_norm 629 if self.override_lipschitz_pair is not None:
618 630 return self.override_lipschitz_pair
619 def diff_bound_pair(self, xbound=None, Delta_t=None): 631
620 if xbound is None: 632 C_mu_adj, C_k_adj, C_c_adj = self.diff_bound3(Delta_t=Delta_t)
621 return self.diff_bound(xbound=None) 633
622 else: 634 if Delta_t is None:
623 xb_combined = xbound[0] + xbound[1] if len(xbound) > 1 else xbound[0] 635 Delta_t = self.dt
624 636
625 k_min = getattr(xb_combined, "k_min", self.alpha) 637 N = int(np.ceil(self.T / Delta_t))
626 c_Linf = getattr(xb_combined, "c_Linf", np.sqrt(self.beta * k_min)) 638
627 diam = getattr(xb_combined, "diam", self.diam) 639 # Extract per-component per-step bounds (smaller scaling)
628 T = getattr(xb_combined, "T", self.T) 640 C_mu_step = C_mu_adj / N
629 641 C_aux_step = max(C_k_adj, C_c_adj) / N
630 Pe = c_Linf * diam / k_min 642
631 gamma = c_Linf**2 / k_min 643 # TINY Lipschitz factors
632 # adaptive time step size Delta_t (Delta_t = 0.01 in the test) 644 C_Lip_mu = C_mu_step * N
633 if Delta_t is None: 645 C_Lip_aux = C_aux_step * N
634 Delta_t = min(0.1, k_min / c_Linf**2, 0.01 * T) # CFL-like 646
635 647 # Light coercivity scaling
636 # Accumulate N = T/Delta_t local steps 648 alpha = self.alpha
637 N = int(np.ceil(T / Delta_t)) 649 inv_alpha = 1.0 / alpha
638 650
639 # Per-step factors 651 C_mu_tot = C_Lip_mu * inv_alpha
640 exp_local = np.exp(0.5 * gamma * Delta_t) 652 C_aux_tot = C_Lip_aux * inv_alpha
641 C_adj_step = np.sqrt(Delta_t * (exp_local + Delta_t / k_min)) 653
642 C_state_step = np.sqrt(Delta_t) * np.sqrt(1 + Pe**2) 654 return C_mu_tot, C_aux_tot
643 655
644 # Total accumulated bound (geometric sum <= N * max_step) 656 def diff_bound3(self, Delta_t=None):
645 C_adj_PDE = N * C_adj_step 657 xbound = self.xbound
646 C_state = N * C_state_step 658 xb_combined = xbound # xbound[0] + xbound[1] if len(xbound) > 1 else xbound[0]
647 659
648 # Use the embedding constant 660 k_min = getattr(xb_combined, "k_min", self.alpha)
649 C_mu_adj = self.C_emb * C_adj_PDE # ∫φ δμ 661 c_Linf = getattr(xb_combined, "c_Linf", np.sqrt(self.beta * k_min))
650 C_k_adj = C_state * C_adj_PDE # ∫∇u·∇φ δk 662 diam = getattr(xb_combined, "diam", self.diam)
651 C_c1_adj = C_state * C_adj_PDE * c_Linf # ∫∂x u φ δc1 663 T = getattr(xb_combined, "T", self.T)
652 C_c2_adj = C_state * C_adj_PDE * c_Linf # ∫∂y u φ δc2 664
653 665 Pe = c_Linf * diam / k_min
654 return C_mu_adj, C_k_adj, C_c1_adj, C_c2_adj 666 gamma = c_Linf**2 / k_min
667 # adaptive time step size Delta_t (Delta_t = 0.01 in the test)
668 if Delta_t is None:
669 Delta_t = min(0.1, k_min / c_Linf**2, 0.01 * T) # CFL-like
670
671 # Accumulate N = T/Delta_t local steps
672 N = int(np.ceil(T / Delta_t))
673
674 # Per-step factors
675 exp_local = np.exp(0.5 * gamma * Delta_t)
676 C_adj_step = np.sqrt(Delta_t * (exp_local + Delta_t / k_min))
677 C_state_step = np.sqrt(Delta_t) * np.sqrt(1 + Pe**2)
678
679 # Total accumulated bound (geometric sum <= N * max_step)
680 C_adj_PDE = N * C_adj_step
681 C_state = N * C_state_step
682
683 # Use the embedding constant
684 C_mu_adj = self.C_emb * C_adj_PDE # ∫φ δμ
685 C_k_adj = C_state * C_adj_PDE # ∫∇u·∇φ δk
686 C_c_adj = C_state * C_adj_PDE * c_Linf # ∫∂x u φ δc1 and c2
687
688 return C_mu_adj, C_k_adj, C_c_adj
655 689
656 # ||S(x)||_Y→X ≤ C_state ||μ||_M(Ω) [time-independent μ] 690 # ||S(x)||_Y→X ≤ C_state ||μ||_M(Ω) [time-independent μ]
657 def codomain_bound(self, xbound=None): 691 def codomain_bound(self):
658 if xbound is None: 692 xbound = self.xbound
659 return 1.0
660 693
661 # Extract parameters 694 # Extract parameters
662 if hasattr(xbound, "k_min"): 695 if hasattr(xbound, "k_min"):
663 xb = xbound 696 xb = xbound
664 else: 697 elif hasattr(xbound, "__len__") and len(xbound) > 0:
665 if hasattr(xbound, "__len__") and len(xbound) > 0: 698 xb = xbound[0] + xbound[1] if len(xbound) > 1 else xbound[0]
666 xb = xbound[0] + xbound[1] if len(xbound) > 1 else xbound[0] 699 else:
667 else: 700 xb = xbound
668 xb = xbound 701
669 k_min = getattr(xb, "k_min", self.alpha) 702 k_min = getattr(xb, "k_min", self.alpha)
670 c_Linf = getattr(xb, "c_Linf", np.sqrt(self.beta * k_min)) 703 c_Linf = getattr(xb, "c_Linf", np.sqrt(self.beta * k_min))
671 mu_dual = getattr(xb, "mu_dual", 3.0) 704 mu_dual = getattr(xb, "mu_dual", 3.0)
672 diam = getattr(xb, "diam", self.diam) 705 diam = getattr(xb, "diam", self.diam)
673 T = getattr(xb, "T", self.T) 706 T = getattr(xb, "T", self.T)
674 707
675 Pe = c_Linf * diam / k_min 708 Pe = c_Linf * diam / k_min
676 if not hasattr(self, "Delta_t"): 709 if not hasattr(self, "Delta_t"):
677 Delta_t = min(0.1, k_min / c_Linf**2, 0.01 * T) 710 Delta_t = min(0.1, k_min / c_Linf**2, 0.01 * T)
678 else: 711 else:
684 717
685 C_state_step = np.sqrt(Delta_t) * np.sqrt(1 + 0.5 * Pe**2) * exp_local 718 C_state_step = np.sqrt(Delta_t) * np.sqrt(1 + 0.5 * Pe**2) * exp_local
686 C_state = N * C_state_step 719 C_state = N * C_state_step
687 720
688 return C_state * mu_dual 721 return C_state * mu_dual
689
690 # ||S(x1)-S(x2)|| ≤ Cμ||Δμ|| + Ck||Δk|| + Cc1||Δc1|| + Cc2||Δc2||
691
692 def codomain_bound_pair(self, xbound=None):
693 if xbound is None:
694 return self.codomain_bound(xbound=None)
695 else:
696 xb = xbound[0] + xbound[1] if len(xbound) > 1 else xbound[0]
697 C_base = self.codomain_bound(xb)
698 # Component
699 C_mu = C_base
700 C_k = C_base * 2.0
701 C_c1 = C_base * 1.5
702 C_c2 = C_base * 1.5
703
704 return (C_mu, C_k, C_c1, C_c2)
705
706 """
707 def codomain_bound(self, xbound=None):
708 return 1.0
709
710 def codomain_bound_pair(self, xbound=None):
711 # TODO: implement properly
712 if xbound is None:
713 return self.codomain_bound(xbound=None)
714 else:
715 return self.codomain_bound(xbound=xbound[0] + xbound[1])
716 """
717
718 def diff_bound3(self):
719 """
720 Returns full operator differential bounds:
721
722 L_e_mu = C_emb * C_reg * sqrt(T)
723 L_e_k = ||∇u||_{L^2(Ω)}
724 L_e_c = ||∇u||_{L^2(Ω)}
725
726 """
727
728 # First get L_e_k = L_e_c
729 grad_u_norm = self.diff_bound()
730
731 # Compute L_e_mu from stability prefactor
732 _P, L_e_mu = self._stability_prefactor()
733
734 L_e_k = grad_u_norm
735 L_e_c = grad_u_norm
736
737 return L_e_mu, L_e_k, L_e_c
738
739 # Solution operator Lipschitz factor separately wrt. μ and (k, c)
740 def lipschitz_factor_pair(self):
741 """
742 Compute the forward solution operator Lipschitz factors
743 separately with respect to the parameters (k, c) and μ.
744
745 Returns:
746 tuple: (L_mu, L_kc)
747 L_mu - Lipschitz factor w.r.t. μ
748 L_kc - Lipschitz factor w.r.t. (k, c)
749 """
750 # Get differential bounds for the solution operator
751 L_e_mu, L_e_k, L_e_c = self.diff_bound3()
752
753 # Lipschitz factor w.r.t μ
754 L_mu = L_e_mu
755
756 # Lipschitz factor w.r.t (k, c)
757 # Conservative choice: max of L_e_k and L_e_c
758 L_kc = max(L_e_k, L_e_c)
759
760 return L_mu, L_kc
761
762 # Adjoint solution operator Lipschitz factor separately wrt. μ and (k, c)
763 def diff_adj_lipschitz_factor_pair(self):
764 A = self.diff_adj_lipschitz_factor()
765
766 # No separate analytic dependence, so return same A for both parts
767 A_mu = A
768 A_kc = A
769
770 return A_mu, A_kc
771 722
772 723
773 def own(u): 724 def own(u):
774 # """ 725 # """
775 # Return a version of u that is guaranteed to be uniquely owned. 726 # Return a version of u that is guaranteed to be uniquely owned.
851 def _get_base(self): 802 def _get_base(self):
852 return (None, (None, None)) if self.base is None else self.base 803 return (None, (None, None)) if self.base is None else self.base
853 804
854 def apply(self, x): 805 def apply(self, x):
855 """ 806 """
856 Apply the quadratic regularisation fucntional to `x`. 807 Apply the quadratic regularisation functional to `x`.
857 """ 808 """
858 (k, (c1, c2)) = x 809 (k, (c1, c2)) = x
859 (kb, (c1b, c2b)) = self._get_base() 810 (kb, (c1b, c2b)) = self._get_base()
860 (nkb, (nc1b, nc2b)) = self.base_norm_squared 811 (nkb, (nc1b, nc2b)) = self.base_norm_squared
861 812
862 return self.α * ( 813 if isinstance(self.α, float):
863 self._apply1(k, kb, nkb) 814 αk, αc = self.α, self.α
864 + self._apply1(c1, c1b, nc1b) 815 else:
865 + self._apply1(c2, c2b, nc2b) 816 αk, αc = self.α
866 ) 817
867 818 return αk * self._apply1(k, kb, nkb) + αc * (
868 def _prox1(self, τ, x, b): 819 self._apply1(c1, c1b, nc1b) + self._apply1(c2, c2b, nc2b)
869 γ = self.α * τ 820 )
821
822 def _prox1(self, γ, x, b):
870 β = 1.0 / (1.0 + γ) 823 β = 1.0 / (1.0 + γ)
871 if isinstance(x, float): 824 if isinstance(x, float):
872 return β * (x if b is None else x + b * γ) 825 return β * (x if b is None else x + b * γ)
873 else: 826 else:
874 # x = own(x) 827 # x = own(x)
878 vx *= β 831 vx *= β
879 return x 832 return x
880 833
881 def prox(self, τ, x): 834 def prox(self, τ, x):
882 """ 835 """
883 Apply the proximal map of the quadratic regularisation fucntional to `x`. 836 Apply the proximal map of the quadratic regularisation functional to `x`.
884 WARNING: This function is unsafe. It may modify `x´. 837 WARNING: This function is unsafe. It may modify `x´.
885 That is ok for our purposes, as Rust, being a safe language, already needs to pass a 838 That is ok for our purposes, as Rust, being a safe language, already needs to pass a
886 copied or owned instance to Python. 839 copied or owned instance to Python.
887 """ 840 """
888 k, (c1, c2) = x 841 k, (c1, c2) = x
889 (kb, (c1b, c2b)) = self._get_base() 842 (kb, (c1b, c2b)) = self._get_base()
890 843
844 if isinstance(self.α, float):
845 αk, αc = self.α, self.α
846 else:
847 αk, αc = self.α
848
891 return ( 849 return (
892 self._prox1(τ, k, kb), 850 self._prox1(τ * αk, k, kb),
893 (self._prox1(τ, c1, c1b), self._prox1(τ, c2, c2b)), 851 (self._prox1(τ * αc, c1, c1b), self._prox1(τ * αc, c2, c2b)),
894 ) 852 )
853
854
855 class LogPowerRegularisation:
856 """
857 Logarithmic barrier + quadratic regularisation functional for the auxiliary variables;
858 $-α\\log(x) + ωx^2$.
859 """
860
861 def _own(self, x):
862 return x if isinstance(x, float) else own(x)
863
864 def __init__(self, α, ω):
865 self.α = α
866 self.ω = ω
867
868 def _apply1(self, x, α, ω):
869 if isinstance(x, float):
870 return -α * np.log(x) + ω * x**2
871 else:
872 # Numpy is Matlab-grade junk. No reduce! Temporary arrays! In 2026! WTF?!?!?!?
873 return (-α * np.log(x) + ω * x**2).sum()
874
875 def apply(self, x):
876 (k, (c1, c2)) = x
877
878 if isinstance(self.α, float):
879 αk, αc = self.α, self.α
880 else:
881 αk, αc = self.α
882 if isinstance(self.ω, float):
883 ωk, ωc = self.ω, self.ω
884 else:
885 ωk, ωc = self.ω
886
887 return (
888 self._apply1(k, αk, ωk)
889 + self._apply1(c1, αc, ωc)
890 + self._apply1(c2, αc, ωc)
891 )
892
893 def _prox1(self, α, ω, x):
894 # OC: 0 = -α/z + 2ωz + z-x ⟺ -α + (2ω+1)z^2 - xz = 0.
895 # ⟺ z = -x ± √(x^2 + 4α(2ω+1))/(2(2ω+1))
896 # Numpy is Matlab-grade junk. No reduce! Temporary arrays! In 2026! WTF?!?!?!?
897 return (x + np.sqrt(x**2 + 4 * α * (2 * ω + 1))) / (2 * (2 * ω + 1))
898
899 def prox(self, τ, x):
900 """
901 Apply the proximal map of the logarithmic regularisation functional to `x`.
902 WARNING: This function is unsafe. It may modify `x´.
903 That is ok for our purposes, as Rust, being a safe language, already needs to pass a
904 copied or owned instance to Python.
905 """
906 k, (c1, c2) = x
907
908 if isinstance(self.α, float):
909 αk, αc = self.α, self.α
910 else:
911 αk, αc = self.α
912 if isinstance(self.ω, float):
913 ωk, ωc = self.ω, self.ω
914 else:
915 ωk, ωc = self.ω
916
917 return (
918 self._prox1(τ * αk, τ * ωk, k),
919 (self._prox1(τ * αc, τ * ωc, c1), self._prox1(τ * αc, τ * ωc, c2)),
920 )
921
922
923 class BoxConstraints:
924 """
925 Simple box constraints on the auxiliary variables
926 """
927
928 def _own(self, x):
929 return x if isinstance(x, float) else own(x)
930
931 def __init__(self, α, ω):
932 self.α = α
933 self.ω = ω
934
935 def _apply1(self, x, α, ω):
936 if isinstance(x, float):
937 if α <= x and x <= ω:
938 return 0.0
939 else:
940 return np.inf
941 else:
942 raise Exception("Unimplemented")
943
944 def apply(self, x):
945 (k, (c1, c2)) = x
946
947 if isinstance(self.α, float):
948 αk, αc = self.α, self.α
949 else:
950 αk, αc = self.α
951 if isinstance(self.ω, float):
952 ωk, ωc = self.ω, self.ω
953 else:
954 ωk, ωc = self.ω
955
956 return (
957 self._apply1(k, αk, ωk)
958 + self._apply1(c1, αc, ωc)
959 + self._apply1(c2, αc, ωc)
960 )
961
962 def _prox1(self, α, ω, x):
963 # OC: 0 = -α/z + 2ωz + z-x ⟺ -α + (2ω+1)z^2 - xz = 0.
964 # ⟺ z = -x ± √(x^2 + 4α(2ω+1))/(2(2ω+1))
965 # Numpy is Matlab-grade junk. No reduce! Temporary arrays! In 2026! WTF?!?!?!?
966 return max(min(ω, x), α)
967
968 def prox(self, τ, x):
969 """
970 Apply the proximal map of the box constraints to `x`.
971 WARNING: This function is unsafe. It may modify `x´.
972 That is ok for our purposes, as Rust, being a safe language, already needs to pass a
973 copied or owned instance to Python.
974 """
975 k, (c1, c2) = x
976
977 if isinstance(self.α, float):
978 αk, αc = self.α, self.α
979 else:
980 αk, αc = self.α
981 if isinstance(self.ω, float):
982 ωk, ωc = self.ω, self.ω
983 else:
984 ωk, ωc = self.ω
985
986 return (
987 self._prox1(αk, ωk, k),
988 (self._prox1(αc, ωc, c1), self._prox1(αc, ωc, c2)),
989 )
990
991
992 """
993 Quadratic regularisation with box contraints
994 """
995
996
997 class BoxedQuadraticRegularisation:
998 def __init__(self, ω0, ω1, α, base=None):
999 self.box = BoxConstraints(ω0, ω1)
1000 self.quadratic = QuadraticRegularisation(α, base)
1001
1002 def apply(self, x):
1003 return self.box.apply(x) + self.quadratic.apply(x)
1004
1005 def prox(self, τ, x):
1006 return self.box.prox(1.0, self.quadratic.prox(τ, x))

Mercurial Repository: pointsource_pde

mercurial