Re-wrote lambda function class so that it can implicitly initialize using a structure constructor (I think this is super cool, even though the minimization method isn't working here yet)

src/modules/lambda_function.f90

@@ -4,33 +4,36 @@ module lambda_function
    implicit none
 
    type, public :: lambda_obj 
-      private
       procedure(lambda0), pointer, nopass :: lambdaptr => null()
+
    contains
       generic   :: init => lambda_init_0
       procedure :: eval => lambda_eval_0
-      procedure :: lambda_init_0
+      procedure, nopass :: lambda_init_0
       final     :: lambda_destroy
    end type
+   interface lambda_obj
+      module procedure lambda_init_0
+   end interface
 
    abstract interface
       function lambda0(x) result(y)
          ! Template for a 0 argument function
          import DP
          real(DP), dimension(:), intent(in) :: x
-         real(DP) :: y
+         real(DP)                           :: y
       end function
    end interface
 
    contains
-      subroutine lambda_init_0(self, lambda)
+      type(lambda_obj) function lambda_init_0(lambda)
          implicit none
          ! Arguments
-         class(lambda_obj), intent(out) :: self
          procedure(lambda0)             :: lambda
 
-         self%lambdaptr => lambda
-      end subroutine lambda_init_0
+         lambda_init_0%lambdaptr => lambda
+         return
+      end function lambda_init_0
 
       function lambda_eval_0(self, x) result(y)
          implicit none

src/modules/module_interfaces.f90

@@ -1579,15 +1579,16 @@ function util_solve_linear_system(A,b,n,lerr) result(x)
             real(DP), dimension(n)               :: x
          end function util_solve_linear_system
 
-         function util_minimize_bfgs(f, N, x1, eps) result(fnum)
+         function util_minimize_bfgs(f, N, x0, eps, lerr) result(x1)
             use swiftest_globals
             use lambda_function
             implicit none
-            integer(I4B), intent(in) :: N
-            class(lambda_obj), intent(in) :: f
-            real(DP), dimension(:), intent(inout) :: x1
-            real(DP), intent(in) :: eps
-            integer(I4B) :: fnum
+            integer(I4B),           intent(in)  :: N
+            class(lambda_obj),      intent(in)  :: f
+            real(DP), dimension(:), intent(in)  :: x0
+            real(DP),               intent(in)  :: eps
+            logical,                intent(out) :: lerr
+            real(DP), dimension(:), allocatable :: x1
          end function util_minimize_bfgs
       end interface
 

src/modules/module_symba.f90

@@ -103,76 +103,76 @@ MODULE module_symba
 
    end type symba_merger 
 
-   type, public, extends(lambda_obj) :: symba_vel_lambda_obj
+   type, public, extends(lambda_obj) :: symba_frag_lambda
       procedure(abstract_objective_func), pointer, nopass :: ke_objective_func_ptr => null()
-      real(DP), dimension(:),   allocatable :: m_frag, v_r_mag
-      real(DP), dimension(:,:), allocatable :: v_r_unit
-      real(DP), dimension(NDIM)             :: L_lin_tan
-      real(DP)                              :: T_rad
+      real(DP), dimension(:),   allocatable :: m_frag
+      real(DP), dimension(:,:), allocatable :: x_frag
+      real(DP), dimension(:,:), allocatable :: v_frag
+      real(DP), dimension(NDIM)             :: L_target
+      real(DP)                              :: ke_target
    contains
       generic   :: init => ke_objective_func_init
       procedure :: eval => ke_objective_func_eval
-      procedure :: ke_objective_func_init
+      procedure, nopass :: ke_objective_func_init
       final     :: ke_objective_func_destroy
-   end type symba_vel_lambda_obj
+   end type symba_frag_lambda
+   interface symba_frag_lambda
+      module procedure ke_objective_func_init
+   end interface
 
    abstract interface
-      function abstract_objective_func(v_r_mag_unknowns, v_r_mag_knowns, m_frag, v_r_unit, L_lin_tan, T_rad) result(fnorm)
+      function abstract_objective_func(vflat, v_frag, x_frag, m_frag, L_target, ke_target) result(fnorm)
          ! Template for the kinetic energy constraint function used for minimizing
          import DP
-         real(DP), dimension(:),   intent(in) :: v_r_mag_unknowns   !! Unknown radial velocity magnitudes
-         real(DP), dimension(:),   intent(in) :: v_r_mag_knowns  !! Known Radial velocity magnitude
-         real(DP), dimension(:),   intent(in) :: m_frag    !! Fragment masses
-         real(DP), dimension(:,:), intent(in) :: v_r_unit  !! Radial unit vectors
-         real(DP), dimension(:),   intent(in) :: L_lin_tan !! Tangential component of linear momentum
-         real(DP),                 intent(in) :: T_rad     !! Target radial kinetic energ
-         real(DP)                             :: fnorm     !! The objective function result: norm of the vector composed of the tangential momentum and energy
+         real(DP), dimension(:),   intent(in) :: vflat          !! Unrolled unknown velocity vectors
+         real(DP), dimension(:,:), intent(in) :: v_frag, x_frag !! Velocity and position vectors
+         real(DP), dimension(:),   intent(in) :: m_frag         !! Fragment masses
+         real(DP), dimension(:),   intent(in) :: L_target       !! Target orbital momentum
+         real(DP),                 intent(in) :: ke_target      !! Target kinetic energ
+         real(DP)                             :: fnorm          !! The objective function result: norm of the vector composed of the tangential momentum and energy
       end function
    end interface
 
    contains
-      subroutine ke_objective_func_init(self, lambda, v_r_mag, m_frag, v_r_unit, L_lin_tan, T_rad)
+      type(symba_frag_lambda) function ke_objective_func_init(lambda, v_frag, x_frag, m_frag, L_target, ke_target)
          implicit none
          ! Arguments
-         class(symba_vel_lambda_obj), intent(out) :: self
          procedure(abstract_objective_func)       :: lambda  !! The lambda function
-         real(DP), dimension(:),      intent(in)  :: v_r_mag   !! Radial velocity magnitude
-         real(DP), dimension(:),      intent(in)  :: m_frag    !! Fragment masses
-         real(DP), dimension(:,:),    intent(in)  :: v_r_unit  !! Radial unit vectors
-         real(DP), dimension(:),      intent(in)  :: L_lin_tan !! Tangential component of linear momentum
-         real(DP),                    intent(in)  :: T_rad     !! Target radial kinetic ener
-
-         self%ke_objective_func_ptr  => lambda
-         allocate(self%m_frag, source=m_frag)
-         allocate(self%v_r_mag, source=v_r_mag)
-         allocate(self%v_r_unit, source=v_r_unit)
-         self%L_lin_tan(:) = L_lin_tan(:)
-         self%T_rad = T_rad
-      end subroutine ke_objective_func_init
+         real(DP), dimension(:,:), intent(in) :: v_frag, x_frag !! Velocity and position vectors
+         real(DP), dimension(:),   intent(in) :: m_frag         !! Fragment masses
+         real(DP), dimension(:),   intent(in) :: L_target       !! Target orbital momentum
+         real(DP),                 intent(in) :: ke_target      !! Target kinetic energ
+         ! Internals
+         associate(self => ke_objective_func_init)
+            self%ke_objective_func_ptr  => lambda
+            allocate(self%m_frag, source=m_frag)
+            allocate(self%v_frag, source=v_frag)
+            allocate(self%x_frag, source=x_frag)
+            self%L_target(:) = L_target(:)
+            self%ke_target = ke_target
+         end associate
+         return
+      end function ke_objective_func_init
 
       subroutine ke_objective_func_destroy(self)
          implicit none
-         type(symba_vel_lambda_obj) :: self
+         type(symba_frag_lambda) :: self
          if (allocated(self%m_frag)) deallocate(self%m_frag)
-         if (allocated(self%v_r_unit)) deallocate(self%v_r_unit)
-         if (allocated(self%v_r_mag)) deallocate(self%v_r_mag)
+         if (allocated(self%v_frag)) deallocate(self%v_frag)
+         if (allocated(self%x_frag)) deallocate(self%x_frag)
          if (associated(self%ke_objective_func_ptr)) nullify(self%ke_objective_func_ptr)
       end subroutine ke_objective_func_destroy 
 
    function ke_objective_func_eval(self, x) result(fnorm)
       implicit none
       ! Arguments
-      class(symba_vel_lambda_obj), intent(in) :: self
-      real(DP), dimension(:),      intent(in) :: x
+      class(symba_frag_lambda), intent(in) :: self
+      real(DP), dimension(:),   intent(in) :: x
       ! Result
       real(DP)                      :: fnorm
-      ! Internals
-      integer(I4B) :: nfrag, nunknown
 
       if (associated(self%ke_objective_func_ptr)) then
-         nfrag = size(self%v_r_mag)
-         nunknown = size(x)
-         fnorm = self%ke_objective_func_ptr(x, self%v_r_mag(nunknown + 1:nfrag), self%m_frag, self%v_r_unit, self%L_lin_tan, self%T_rad)
+         fnorm = self%ke_objective_func_ptr(x, self%v_frag, self%x_frag, self%m_frag, self%L_target, self%ke_target)
       else
          error stop "KE Objective function was not initialized."
       end if

src/symba/symba_frag_pos.f90

@@ -25,7 +25,7 @@ subroutine symba_frag_pos (param, symba_plA, family, x, v, L_spin, Ip, mass, rad
    real(DP), dimension(:,:), allocatable   :: x_frag, v_frag ! Fragment positions and velocities in the collision center of mass frame
    real(DP), dimension(NDIM, 2)            :: rot, L_orb 
    integer(I4B)                            :: i, j, nfrag, fam_size, istart
-   real(DP), dimension(NDIM)               :: xcom, vcom, Ltot_before, Ltot_after, L_residual, L_spin_frag
+   real(DP), dimension(NDIM)               :: xcom, vcom, Ltot_before, Ltot_after, L_residual, L_spin_frag, L_target, L_frag
    real(DP)                                :: mtot, Lmag_before, Lmag_after
    real(DP)                                :: Etot_before, Etot_after, ke_before, pe_before
    real(DP)                                :: pe_after, ke_spin_before, ke_spin_after, ke_after, ke_family, ke_target, ke_frag
@@ -115,7 +115,12 @@ subroutine symba_frag_pos (param, symba_plA, family, x, v, L_spin, Ip, mass, rad
 
       ! Set the "target" ke_after (the value of the orbital kinetic energy that the fragments ought to have)
       ke_target = ke_family + (ke_spin_before - ke_spin_after) + (pe_before - pe_after) - Qloss
-      call symba_frag_pos_kinetic_energy(xcom, vcom, L_orb, L_spin, m_frag, x_frag, v_frag, ke_target, lmerge)
+      L_target(:) = 0.0_DP
+      do i = 1, nfrag
+         call util_crossproduct(x_frag(:, i), v_frag(:, i), L_frag(:))
+         L_target(:) = L_target(:) + L_frag(:)
+      end do
+      call symba_frag_pos_kinetic_energy(m_frag, ke_target, L_target, x_frag, v_frag, lmerge)
 
       write(*,        "(' ---------------------------------------------------------------------------')")
       write(*,fmtlabel) ' T_family    |',ke_family / abs(Etot_before)
@@ -320,7 +325,7 @@ subroutine symba_frag_pos_ang_mtm(nfrag, xcom, vcom, L_orb, L_spin, m_frag, rad_
       return 
    end subroutine symba_frag_pos_ang_mtm
 
-   subroutine symba_frag_pos_kinetic_energy(xcom, vcom, L_orb, L_spin, m_frag, x_frag, v_frag, ke_target, lmerge)
+   subroutine symba_frag_pos_kinetic_energy(m_frag, ke_target, L_target, x_frag, v_frag, lmerge)
       !! Author: Jennifer L.L. Pouplin, Carlisle A. Wishard, and David A. Minton
       !!
       !! 
@@ -332,96 +337,82 @@ subroutine symba_frag_pos_kinetic_energy(xcom, vcom, L_orb, L_spin, m_frag, x_fr
       !! the target kinetic energy required to satisfy the constraints.
       implicit none
       ! Arguments
-      real(DP), dimension(:),   intent(in)    :: xcom, vcom      !! Center of mass position and velocity in the system barycenter frame
-      real(DP), dimension(:,:), intent(in)    :: L_orb, L_spin   !! Pre-impact orbital and spin angular momentum
-      real(DP), dimension(:),   intent(in)    :: m_frag          !! Fragment masses
-      real(DP), dimension(:,:), intent(inout) :: x_frag, v_frag  !! Fragment position and velocities in the center of mass frame   
-      real(DP), intent(in)                    :: ke_target        !! Target kinetic energy 
-      logical, intent(out)                    :: lmerge
+      real(DP), dimension(:),   intent(in)    :: m_frag     !! Fragment masses
+      real(DP), dimension(:,:), intent(in)    :: x_frag     !! Fragment position vectors
+      real(DP),                 intent(in)    :: ke_target  !! Target kinetic energy 
+      real(DP), dimension(:),   intent(in)    :: L_target   !! Target kinetic energy 
+      real(DP), dimension(:,:), intent(inout) :: v_frag      !! Fragment position and velocities in the center of mass frame   
+      logical,                  intent(out)   :: lmerge
 
       ! Internals
       real(DP)                              :: mtot           !! Total mass of fragments
-      real(DP)                              :: T_rad          !! Sum of the radial kinetic energy of all fragments 
-      real(DP), dimension(NDIM)             :: L_lin_tan            !! Sum of the tangential momentum vector of all fragments
       integer(I4B)                          :: i, nfrag, neval
-      real(DP), dimension(:,:), allocatable :: v_r_unit, v_t
-      real(DP), dimension(NDIM)             :: v_t_unit, h_unit, L_orb_frag
-      real(DP), dimension(:), allocatable   :: v_r_mag, v_r_mag_sol
-      integer(I4B), parameter               :: MAXITER = 500 
-      real(DP), parameter                   :: TOL = 1e-12_DP
-      real(DP)                              :: vmult
-      type(symba_vel_lambda_obj)            :: ke_objective_func
+      real(DP), parameter                   :: TOL = 1e-9_DP
+      real(DP), dimension(:), allocatable   :: vflat
+      logical                               :: lerr
+
 
       nfrag = size(m_frag)
       mtot = sum(m_frag)
 
-      allocate(v_r_unit, mold=v_frag) 
-      allocate(v_t, mold=v_frag)
-      allocate(v_r_mag, mold=m_frag)
-
-      ! Create the radial unit vectors pointing away from the collision center of mass, and subtract that off of the current
-      ! fragment velocities in order to create the tangential component 
-      do i = 1, nfrag
-         v_r_unit(:, i) = x_frag(:,i) / norm2(x_frag(:, i))
-         v_t(:,i) = v_frag(:, i) - dot_product(v_frag(:,i), v_r_unit(:, i)) * v_r_unit(:, i)
-      end do
-
-      L_lin_tan = 0.0_DP
-      T_rad = ke_target - 0.5_DP * mtot * dot_product(vcom(:), vcom(:))
-      do i = 1, nfrag
-         T_rad = T_rad - 0.5_DP * m_frag(i) * dot_product(v_t(:, i), v_t(:, i))
-         L_lin_tan(:) = L_lin_tan(:) + m_frag(i) * v_t(:, i)
-      end do
-      if (T_rad > 0.0_DP) then
-         lmerge = .false.
-
-         call random_number(v_r_mag(:))
-         v_r_mag(:) = sqrt(2 * T_rad / nfrag / m_frag(:)) * (v_r_mag(:) + 0.1_DP)
-
-         ! Initialize the lambda function with all the parameters that stay constant during the minimization
-         call ke_objective_func%init(symba_frag_pos_ke_objective_function, v_r_mag, m_frag, v_r_unit, L_lin_tan, T_rad)
-         ! Minimize error using the BFGS optimizer
-         neval = util_minimize_bfgs(ke_objective_func, 4, v_r_mag(1:4), TOL)
-
-         do i = 1, nfrag
-            v_frag(:, i) = v_r_mag(i) * v_r_unit(:, i) + v_t(:, i)
-         end do
-      else
+      ! Initialize the lambda function using a structure constructor that calls the init method
+      ! Minimize error using the BFGS optimizer
+      vflat = util_minimize_bfgs(symba_frag_lambda(lambda=symba_frag_pos_ke_objective_function, v_frag=v_frag, x_frag=x_frag, m_frag=m_frag, L_target=L_target, ke_target=ke_target), &
+                                 NDIM*nfrag, reshape(v_frag,[NDIM*nfrag]), TOL, lerr)
+      if (lerr) then
          ! No solution exists for this case, so we need to indicate that this should be a merge
          ! This may happen due to setting the tangential velocities too high when setting the angular momentum constraint
          lmerge = .true.
-         v_frag(:, :) = 0.0_DP
+         v_frag(:,:) = 0.0_DP
+      else
+         v_frag(:,:) = reshape(vflat,shape(v_frag))
       end if
 
       return
    end subroutine symba_frag_pos_kinetic_energy
 
-   function symba_frag_pos_ke_objective_function(v_r_mag_unknowns, v_r_mag_knowns, m_frag, v_r_unit, L_lin_tan, T_rad) result(fnorm)
+   function symba_frag_pos_ke_objective_function(vflat, v_frag, x_frag, m_frag, L_target, ke_target) result(fnorm)
       ! Objective function for evaluating how close our fragment velocities get to minimizing KE error from our required value
       implicit none
       ! Arguments
-      real(DP), dimension(:),   intent(in) :: v_r_mag_unknowns !! Radial velocity magnitude
-      real(DP), dimension(:),   intent(in) :: v_r_mag_knowns   !! Radial velocity magnitude
-      real(DP), dimension(:),   intent(in) :: m_frag    !! Fragment masses
-      real(DP), dimension(:),   intent(in) :: L_lin_tan !! Tangential component of linear momentum
-      real(DP), dimension(:,:), intent(in) :: v_r_unit  !! Radial unit vectors
-      real(DP),                 intent(in) :: T_rad     !! Target radial kinetic energy
+      real(DP), dimension(:),   intent(in) :: vflat          !! Unrolled unknown velocity vectors
+      real(DP), dimension(:,:), intent(in) :: v_frag, x_frag !! Velocity and position vectors
+      real(DP), dimension(:),   intent(in) :: m_frag         !! Fragment masses
+      real(DP), dimension(:),   intent(in) :: L_target       !! Target orbital momentum
+      real(DP),                 intent(in) :: ke_target      !! Target kinetic energ
       ! Result
       real(DP)                             :: fnorm     !! The objective function result: norm of the vector composed of the tangential momentum and energy
                                                         !! Minimizing this brings us closer to our objective
       ! Internals
-      real(DP), dimension(4)  :: f_vec, f_vec_0
-
-      f_vec_0(1:3) = L_lin_tan(1:3)
-      f_vec_0(4) = -T_rad
-      f_vec(1) = sum(m_frag(:) * [v_r_mag_unknowns(:), v_r_mag_knowns(:)] * v_r_unit(1,:))
-      f_vec(2) = sum(m_frag(:) * [v_r_mag_unknowns(:), v_r_mag_knowns(:)] * v_r_unit(2,:))
-      f_vec(3) = sum(m_frag(:) * [v_r_mag_unknowns(:), v_r_mag_knowns(:)] * v_r_unit(3,:))
-      f_vec(4) = sum(m_frag(:) * [v_r_mag_unknowns(:), v_r_mag_knowns(:)] **2) 
+      real(DP), dimension(7)  :: f_vec
+      integer(I4B) :: i, nfrag, nsol
+      real(DP), dimension(NDIM) :: L
+      real(DP), dimension(:,:), allocatable :: v_sol
 
-      f_vec(:) = f_vec(:) + f_vec_0(:)
+      nfrag = size(m_frag)
+      nsol = size(vflat) / NDIM
+      allocate(v_sol(NDIM,nsol))
+      v_sol = reshape(vflat, shape(v_sol))
+
+      ! Linear momentum constraint
+      f_vec(1) = sum(m_frag(:) * [v_sol(1, 1:nsol), v_frag(1, nsol + 1:nfrag)])
+      f_vec(2) = sum(m_frag(:) * [v_sol(2, 1:nsol), v_frag(2, nsol + 1:nfrag)])
+      f_vec(3) = sum(m_frag(:) * [v_sol(3, 1:nsol), v_frag(3, nsol + 1:nfrag)])
+      ! Angular momentum and kinetic energy constraints
+      f_vec(4:6) = -L_target(:)
+      f_vec(7) = -ke_target
+      do i = 1, nsol
+         call util_crossproduct(x_frag(:,i), v_sol(:, i), L(:))
+         f_vec(4:6) = f_vec(4:6) + m_frag(i) * L(:)
+         f_vec(7)   = f_vec(7)   + 0.5_DP * m_frag(i) * dot_product(v_sol(:,i), v_sol(:, i))
+      end do
+      do i = nsol + 1, nfrag
+         call util_crossproduct(x_frag(:,i), v_frag(:, i), L(:))
+         f_vec(4:6) = f_vec(4:6) + m_frag(i) * L(:)
+         f_vec(7)   = f_vec(7) + 0.5_DP * m_frag(i) * dot_product(v_frag(:,i), v_frag(:, i))
+      end do
 
-      fnorm = norm2(f_vec(:) / f_vec_0(:))
+      fnorm = norm2(f_vec(:)) 
 
       return