From 32f44c5cba0119ed82f6ff57c37d5f5e6c91a7b7 Mon Sep 17 00:00:00 2001
From: Michele Guerini Rocco <m.guerinirocco2@campus.unimib.it>
Date: Wed, 31 Jan 2024 17:10:15 +0100
Subject: [PATCH] src/splines.f90: add procedure to transform spline_2d

This adds a proper procedure to rescale and shift a 2D B-spline by
manipulating the coefficients.

Note: the previous code in set_equil_spline did work, but by
transforming the whole partial(i,j) triggered warnings about operations
on uninitialised memory.
---
 src/equilibrium.f90 | 13 ++-----------
 src/splines.f90     | 39 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 41 insertions(+), 11 deletions(-)

diff --git a/src/equilibrium.f90 b/src/equilibrium.f90
index a25dfb1..ad8cb61 100644
--- a/src/equilibrium.f90
+++ b/src/equilibrium.f90
@@ -647,17 +647,8 @@ contains
         call log_info(msg, mod='equilibrium', proc='set_equil_spline')
     end select
 
-    ! Adjust all the B-spline coefficients
-    ! Note: since ψ_n(R,z) = Σ_ij c_ij B_i(R)B_j(z), to correct ψ_n
-    ! it's enough to correct (shift and scale) the c_ij
-    psi_spline%coeffs  = (psi_spline%coeffs - psinop) / psiant
-
-    ! Same for the derivatives
-    psi_spline%partial(1,0)%ptr = psi_spline%partial(1,0)%ptr / psiant
-    psi_spline%partial(0,1)%ptr = psi_spline%partial(0,1)%ptr / psiant
-    psi_spline%partial(2,0)%ptr = psi_spline%partial(2,0)%ptr / psiant
-    psi_spline%partial(0,2)%ptr = psi_spline%partial(0,2)%ptr / psiant
-    psi_spline%partial(1,1)%ptr = psi_spline%partial(1,1)%ptr / psiant
+    ! Correct the spline coefficients: ψ_n → (ψ_n - psinop)/psiant
+    call psi_spline%transform(1/psiant, -psinop/psiant)
 
     ! Do the opposite scaling to preserve un-normalised values
     ! Note: this is only used for the poloidal magnetic field
diff --git a/src/splines.f90 b/src/splines.f90
index 2e30f0a..99f88ab 100644
--- a/src/splines.f90
+++ b/src/splines.f90
@@ -51,6 +51,7 @@ module splines
       procedure :: eval       => spline_2d_eval
       procedure :: init_deriv => spline_2d_init_deriv
       procedure :: deriv      => spline_2d_deriv
+      procedure :: transform  => spline_2d_transform
   end type
 
   ! Wrapper to store pointers in an array
@@ -532,6 +533,44 @@ contains
   end function spline_2d_deriv
 
 
+  subroutine spline_2d_transform(self, a, b)
+    ! Applies an affine transformation to the spline by
+    ! a manipulation of its B-spline coefficients.
+    ! If the spline is s(x,y), this maps s(x,y) → a⋅s(x,y) + b
+
+    ! subroutine arguments
+    class(spline_2d), intent(inout) :: self
+    real(wp_),        intent(in)    :: a, b
+
+    ! local variables
+    integer :: i, j, m
+
+    ! Note: since s(x,y) = Σ_ij c_ij B_i(x)B_j(y), the affine
+    ! transformation simply acts on the spline coefficients as
+    ! c_ij → a⋅c_ij + b
+
+    ! Transform the spline
+    self%coeffs = a * self%coeffs + b
+
+    ! Transform the derivative splines too
+    if (allocated(self%partial)) then
+
+      ! Note: self%partial is a workspace array, only the
+      ! first m elements contain the spline coefficients
+      m = (self%nknots_x - 3 - 1) * (self%nknots_y - 3 - 1)
+
+      do i = 0, size(self%partial, dim=1) - 1
+      do j = 0, size(self%partial, dim=2) - 1
+        if (associated(self%partial(i, j)%ptr)) then
+          self%partial(i, j)%ptr(1:m) = a * self%partial(i, j)%ptr(1:m)
+        end if
+      end do
+      end do
+
+    end if
+  end subroutine spline_2d_transform
+
+
   subroutine linear_1d_init(self, x, y, n)
     ! Initialises the spline