diff --git a/Makefile b/Makefile index 162129149..94dfbceeb 100644 --- a/Makefile +++ b/Makefile @@ -45,13 +45,13 @@ #****************************************************************************** SWIFTEST_MODULES = swiftest_globals.f90 \ + lambda_function.f90\ swiftest_classes.f90 \ swiftest_operators.f90 \ whm_classes.f90 \ rmvs_classes.f90 \ helio_classes.f90 \ symba_classes.f90 \ - lambda_function.f90\ swiftest.f90 diff --git a/Makefile.Defines b/Makefile.Defines index 291f2c604..70069bb71 100644 --- a/Makefile.Defines +++ b/Makefile.Defines @@ -65,13 +65,13 @@ GPAR = -fopenmp -ftree-parallelize-loops=4 GMEM = -fsanitize=undefined -fsanitize=address -fsanitize=leak GWARNINGS = -Wall -Warray-bounds -Wimplicit-interface -Wextra -Warray-temporaries -FFLAGS = $(IDEBUG) $(HEAPARR) +#FFLAGS = $(IDEBUG) $(HEAPARR) #FFLAGS = -init=snan,arrays -no-wrap-margin -O3 $(STRICTREAL) $(SIMDVEC) $(PAR) -FORTRAN = ifort +#FORTRAN = ifort #AR = xiar -#FORTRAN = gfortran -#FFLAGS = -ffree-line-length-none $(GDEBUG) $(GMEM) +FORTRAN = gfortran +FFLAGS = -ffree-line-length-none $(GDEBUG) $(GMEM) AR = ar # DO NOT include in CFLAGS the "-c" option to compile object only diff --git a/src/fragmentation/fragmentation.f90 b/src/fragmentation/fragmentation.f90 index a01af206b..d1965914f 100644 --- a/src/fragmentation/fragmentation.f90 +++ b/src/fragmentation/fragmentation.f90 @@ -1,6 +1,841 @@ submodule(swiftest_classes) s_fragmentation use swiftest contains + + subroutine fragmentation_initialize(system, param, family, x, v, L_spin, Ip, mass, radius, & + nfrag, Ip_frag, m_frag, rad_frag, xb_frag, vb_frag, rot_frag, Qloss, lfailure) + !! Author: Jennifer L.L. Pouplin, Carlisle A. Wishard, and David A. Minton + !! + !! Initialize the position and velocity of fragments to conserve energy and momentum. + use, intrinsic :: ieee_exceptions + implicit none + ! Arguments + class(swiftest_nbody_system), intent(inout) :: system + class(swiftest_parameters), intent(in) :: param + integer(I4B), dimension(:), intent(in) :: family + real(DP), dimension(:,:), intent(inout) :: x, v, L_spin, Ip + real(DP), dimension(:), intent(inout) :: mass, radius + integer(I4B), intent(inout) :: nfrag + real(DP), dimension(:), allocatable, intent(inout) :: m_frag, rad_frag + real(DP), dimension(:,:), allocatable, intent(inout) :: Ip_frag + real(DP), dimension(:,:), allocatable, intent(inout) :: xb_frag, vb_frag, rot_frag + logical, intent(out) :: lfailure ! Answers the question: Should this have been a merger instead? + real(DP), intent(inout) :: Qloss + ! Internals + real(DP) :: mscale, rscale, vscale, tscale, Lscale, Escale ! Scale factors that reduce quantities to O(~1) in the collisional system + real(DP) :: mtot + real(DP), dimension(NDIM) :: xcom, vcom + integer(I4B) :: ii + logical, dimension(:), allocatable :: lexclude + real(DP), dimension(NDIM, 2) :: rot, L_orb + real(DP), dimension(:,:), allocatable :: x_frag, v_frag, v_r_unit, v_t_unit, v_h_unit + real(DP), dimension(:), allocatable :: rmag, rotmag, v_r_mag, v_t_mag + real(DP), dimension(NDIM) :: Ltot_before + real(DP), dimension(NDIM) :: Ltot_after + real(DP) :: Etot_before, ke_orbit_before, ke_spin_before, pe_before, Lmag_before + real(DP) :: Etot_after, ke_orbit_after, ke_spin_after, pe_after, Lmag_after, dEtot, dLmag + real(DP), dimension(NDIM) :: L_frag_tot, L_frag_orb + real(DP) :: ke_frag_budget, ke_frag_orbit, ke_radial, ke_frag_spin, ke_avg_deficit, ke_avg_deficit_old + real(DP), dimension(NDIM) :: x_col_unit, y_col_unit, z_col_unit + character(len=*), parameter :: fmtlabel = "(A14,10(ES11.4,1X,:))" + integer(I4B) :: try, subtry + integer(I4B), parameter :: NFRAG_MIN = 7 !! The minimum allowable number of fragments (set to 6 because that's how many unknowns are needed in the tangential velocity calculation) + real(DP) :: r_max_start, r_max_start_old, r_max, f_spin + real(DP), parameter :: Ltol = 10 * epsilon(1.0_DP) + real(DP), parameter :: Etol = 1e-10_DP + integer(I4B), parameter :: MAXTRY = 3000 + integer(I4B), parameter :: TANTRY = 3 + logical, dimension(size(IEEE_ALL)) :: fpe_halting_modes, fpe_quiet_modes + + if (nfrag < NFRAG_MIN) then + write(*,*) "symba_frag_pos needs at least ",NFRAG_MIN," fragments, but only ",nfrag," were given." + lfailure = .true. + return + end if + + call ieee_get_halting_mode(IEEE_ALL,fpe_halting_modes) ! Save the current halting modes so we can turn them off temporarily + fpe_quiet_modes(:) = .false. + call ieee_set_halting_mode(IEEE_ALL,fpe_quiet_modes) + + f_spin = 0.05_DP + mscale = 1.0_DP + rscale = 1.0_DP + vscale = 1.0_DP + tscale = 1.0_DP + Lscale = 1.0_DP + Escale = 1.0_DP + + associate(npl => system%pl%nbody, status => system%pl%status) + allocate(lexclude(npl)) + where (status(1:npl) == INACTIVE) ! Safety check in case one of the included bodies has been previously deactivated + lexclude(1:npl) = .true. + elsewhere + lexclude(1:npl) = .false. + end where + end associate + + allocate(x_frag, source=xb_frag) + allocate(v_frag, source=vb_frag) + + call set_scale_factors() + call define_coordinate_system() + call calculate_system_energy(linclude_fragments=.false.) + + r_max_start = norm2(x(:,2) - x(:,1)) + try = 1 + lfailure = .false. + ke_avg_deficit = 0.0_DP + do while (try < MAXTRY) + lfailure = .false. + ke_avg_deficit_old = ke_avg_deficit + ke_avg_deficit = 0.0_DP + subtry = 1 + do + call set_fragment_position_vectors() + call set_fragment_tan_vel(lfailure) + ke_avg_deficit = ke_avg_deficit - ke_radial + subtry = subtry + 1 + if (.not.lfailure .or. subtry == TANTRY) exit + !write(*,*) 'Trying new arrangement' + end do + ke_avg_deficit = ke_avg_deficit / subtry + if (lfailure) write(*,*) 'Failed to find tangential velocities' + + if (.not.lfailure) then + call set_fragment_radial_velocities(lfailure) + if (lfailure) write(*,*) 'Failed to find radial velocities' + if (.not.lfailure) then + call calculate_system_energy(linclude_fragments=.true.) + !write(*,*) 'Qloss : ',Qloss + !write(*,*) '-dEtot: ',-dEtot + !write(*,*) 'delta : ',abs((dEtot + Qloss)) + if ((abs(dEtot + Qloss) > Etol) .or. (dEtot > 0.0_DP)) then + write(*,*) 'Failed due to high energy error: ',dEtot, abs(dEtot + Qloss) / Etol + lfailure = .true. + else if (abs(dLmag) / Lmag_before > Ltol) then + write(*,*) 'Failed due to high angular momentum error: ', dLmag / Lmag_before + lfailure = .true. + end if + end if + end if + + if (.not.lfailure) exit + call restructure_failed_fragments() + try = try + 1 + end do + write(*, "(' -------------------------------------------------------------------------------------')") + write(*, "(' Final diagnostic')") + write(*, "(' -------------------------------------------------------------------------------------')") + if (lfailure) then + write(*,*) "symba_frag_pos failed after: ",try," tries" + do ii = 1, nfrag + vb_frag(:, ii) = vcom(:) + end do + else + write(*,*) "symba_frag_pos succeeded after: ",try," tries" + write(*, "(' dL_tot should be very small' )") + write(*,fmtlabel) ' dL_tot |', dLmag / Lmag_before + write(*, "(' dE_tot should be negative and equal to Qloss' )") + write(*,fmtlabel) ' dE_tot |', dEtot / abs(Etot_before) + write(*,fmtlabel) ' Qloss |', -Qloss / abs(Etot_before) + write(*,fmtlabel) ' dE - Qloss |', (Etot_after - Etot_before + Qloss) / abs(Etot_before) + end if + write(*, "(' -------------------------------------------------------------------------------------')") + + call restore_scale_factors() + call ieee_set_halting_mode(IEEE_ALL,fpe_halting_modes) ! Save the current halting modes so we can turn them off temporarily + + return + + contains + + ! Because of the complexity of this procedure, we have chosen to break it up into a series of nested subroutines. + + subroutine set_scale_factors() + !! author: David A. Minton + !! + !! Scales dimenional quantities to ~O(1) with respect to the collisional system. This scaling makes it easier for the non-linear minimization + !! to converge on a solution + implicit none + integer(I4B) :: i + + ! Find the center of mass of the collisional system + mtot = sum(mass(:)) + xcom(:) = (mass(1) * x(:,1) + mass(2) * x(:,2)) / mtot + vcom(:) = (mass(1) * v(:,1) + mass(2) * v(:,2)) / mtot + + ! Set scale factors + !! Because of the implied G, mass is actually G*mass with units of distance**3 / time**2 + Escale = 0.5_DP * (mass(1) * dot_product(v(:,1), v(:,1)) + mass(2) * dot_product(v(:,2), v(:,2))) + rscale = sum(radius(:)) + mscale = sqrt(Escale * rscale) + vscale = sqrt(Escale / mscale) + tscale = rscale / vscale + Lscale = mscale * rscale * vscale + + xcom(:) = xcom(:) / rscale + vcom(:) = vcom(:) / vscale + + mtot = mtot / mscale + mass = mass / mscale + radius = radius / rscale + x = x / rscale + v = v / vscale + L_spin = L_spin / Lscale + do i = 1, 2 + rot(:,i) = L_spin(:,i) / (mass(i) * radius(i)**2 * Ip(3, i)) + end do + + m_frag = m_frag / mscale + rad_frag = rad_frag / rscale + Qloss = Qloss / Escale + + return + end subroutine set_scale_factors + + subroutine restore_scale_factors() + !! author: David A. Minton + !! + !! Restores dimenional quantities back to the system units + implicit none + integer(I4B) :: i + + call ieee_set_halting_mode(IEEE_ALL,.false.) + ! Restore scale factors + xcom(:) = xcom(:) * rscale + vcom(:) = vcom(:) * vscale + + mtot = mtot * mscale + mass = mass * mscale + radius = radius * rscale + x = x * rscale + v = v * vscale + L_spin = L_spin * Lscale + do i = 1, 2 + rot(:,i) = L_spin(:,i) * (mass(i) * radius(i)**2 * Ip(3, i)) + end do + + m_frag = m_frag * mscale + rad_frag = rad_frag * rscale + rot_frag = rot_frag / tscale + x_frag = x_frag * rscale + v_frag = v_frag * vscale + Qloss = Qloss * Escale + + do i = 1, nfrag + xb_frag(:, i) = x_frag(:, i) + xcom(:) + vb_frag(:, i) = v_frag(:, i) + vcom(:) + end do + + Etot_before = Etot_before * Escale + pe_before = pe_before * Escale + ke_spin_before = ke_spin_before * Escale + ke_orbit_before = ke_orbit_before * Escale + Ltot_before = Ltot_before * Lscale + Lmag_before = Lmag_before * Lscale + Etot_after = Etot_after * Escale + pe_after = pe_after * Escale + ke_spin_after = ke_spin_after * Escale + ke_orbit_after = ke_orbit_after * Escale + Ltot_after = Ltot_after * Lscale + Lmag_after = Lmag_after * Lscale + + mscale = 1.0_DP + rscale = 1.0_DP + vscale = 1.0_DP + tscale = 1.0_DP + Lscale = 1.0_DP + Escale = 1.0_DP + + return + end subroutine restore_scale_factors + + subroutine define_coordinate_system() + !! author: David A. Minton + !! + !! Defines the collisional coordinate system, including the unit vectors of both the system and individual fragments. + implicit none + integer(I4B) :: i + real(DP), dimension(NDIM) :: x_cross_v, xc, vc, delta_r, delta_v + real(DP) :: r_col_norm, v_col_norm + + allocate(rmag(nfrag)) + allocate(rotmag(nfrag)) + allocate(v_r_mag(nfrag)) + allocate(v_t_mag(nfrag)) + allocate(v_r_unit(NDIM,nfrag)) + allocate(v_t_unit(NDIM,nfrag)) + allocate(v_h_unit(NDIM,nfrag)) + + rmag(:) = 0.0_DP + rotmag(:) = 0.0_DP + v_r_mag(:) = 0.0_DP + v_t_mag(:) = 0.0_DP + v_r_unit(:,:) = 0.0_DP + v_t_unit(:,:) = 0.0_DP + v_h_unit(:,:) = 0.0_DP + + L_orb(:, :) = 0.0_DP + ! Compute orbital angular momentum of pre-impact system + do i = 1, 2 + xc(:) = x(:, i) - xcom(:) + vc(:) = v(:, i) - vcom(:) + x_cross_v(:) = xc(:) .cross. vc(:) + L_orb(:, i) = mass(i) * x_cross_v(:) + end do + + ! Compute orbital angular momentum of pre-impact system. This will be the normal vector to the collision fragment plane + L_frag_tot(:) = L_spin(:, 1) + L_spin(:, 2) + L_orb(:, 1) + L_orb(:, 2) + + delta_v(:) = v(:, 2) - v(:, 1) + v_col_norm = norm2(delta_v(:)) + delta_r(:) = x(:, 2) - x(:, 1) + r_col_norm = norm2(delta_r(:)) + + ! We will initialize fragments on a plane defined by the pre-impact system, with the z-axis aligned with the angular momentum vector + ! and the y-axis aligned with the pre-impact distance vector. + y_col_unit(:) = delta_r(:) / r_col_norm + z_col_unit(:) = L_frag_tot(:) / norm2(L_frag_tot) + ! The cross product of the y- by z-axis will give us the x-axis + x_col_unit(:) = y_col_unit(:) .cross. z_col_unit(:) + + return + end subroutine define_coordinate_system + + subroutine calculate_system_energy(linclude_fragments) + !! Author: David A. Minton + !! + !! Calculates total system energy, including all bodies in the pl list that do not have a corresponding value of the lexclude array that is true + !! and optionally including fragments. + implicit none + ! Arguments + logical, intent(in) :: linclude_fragments + ! Internals + integer(I4B) :: i, npl_new, nplm + logical, dimension(:), allocatable :: ltmp + logical :: lk_plpl + class(swiftest_nbody_system), allocatable :: tmpsys + + ! Because we're making a copy of symba_pl with the excludes/fragments appended, we need to deallocate the + ! big k_plpl array and recreate it when we're done, otherwise we run the risk of blowing up the memory by + ! allocating two of these ginormous arrays simulteouously. This is not particularly efficient, but as this + ! subroutine should be called relatively infrequently, it shouldn't matter too much. + !if (allocated(pl%k_plpl)) deallocate(pl%k_plpl) + + ! Build the internal planet list out of the non-excluded bodies and optionally with fragments appended. This + ! will get passed to the energy calculation subroutine so that energy is computed exactly the same way is it + ! is in the main program. + associate(pl => system%pl, npl => system%pl%nbody, cb => system%cb) + lk_plpl = allocated(pl%k_plpl) + if (lk_plpl) deallocate(pl%k_plpl) + if (linclude_fragments) then ! Temporarily expand the planet list to feed it into symba_energy + lexclude(family(:)) = .true. + npl_new = npl + nfrag + else + npl_new = npl + end if + call setup_construct_system(tmpsys, param) + deallocate(tmpsys%cb) + allocate(tmpsys%cb, source=cb) + allocate(ltmp(npl)) + ltmp(:) = .true. + call tmpsys%pl%setup(npl_new, param) + call tmpsys%pl%fill(pl, ltmp) + deallocate(ltmp) + + if (linclude_fragments) then ! Append the fragments if they are included + ! Energy calculation requires the fragments to be in the system barcyentric frame, s + tmpsys%pl%mass(npl+1:npl_new) = m_frag(1:nfrag) + tmpsys%pl%radius(npl+1:npl_new) = rad_frag(1:nfrag) + tmpsys%pl%xb(:,npl+1:npl_new) = xb_frag(:,1:nfrag) + tmpsys%pl%vb(:,npl+1:npl_new) = vb_frag(:,1:nfrag) + tmpsys%pl%status(npl+1:npl_new) = COLLISION + if (param%lrotation) then + tmpsys%pl%Ip(:,npl+1:npl_new) = Ip_frag(:,1:nfrag) + tmpsys%pl%rot(:,npl+1:npl_new) = rot_frag(:,1:nfrag) + end if + call tmpsys%pl%b2h(tmpsys%cb) + allocate(ltmp(npl_new)) + ltmp(1:npl) = lexclude(1:npl) + ltmp(npl+1:npl_new) = .false. + call move_alloc(ltmp, lexclude) + end if + + where (lexclude(1:npl_new)) + tmpsys%pl%status(1:npl_new) = INACTIVE + end where + + select type(plwksp => tmpsys%pl) + class is (symba_pl) + select type(param) + class is (symba_parameters) + plwksp%nplm = count(plwksp%Gmass > param%mtiny / mscale) + end select + end select + call tmpsys%pl%eucl_index() + call tmpsys%get_energy_and_momentum(param) + + ! Restore the big array + deallocate(tmpsys%pl%k_plpl) + select type(pl) + class is (symba_pl) + select type(param) + class is (symba_parameters) + nplm = count(pl%mass > param%mtiny) + end select + end select + if (lk_plpl) call pl%eucl_index() + + ! Calculate the current fragment energy and momentum balances + if (linclude_fragments) then + Ltot_after(:) = tmpsys%Lorbit(:) + tmpsys%Lspin(:) + Lmag_after = norm2(Ltot_after(:)) + ke_orbit_after = tmpsys%ke_orbit + ke_spin_after = tmpsys%ke_spin + pe_after = tmpsys%pe + Etot_after = ke_orbit_after + ke_spin_after + pe_after + dEtot = Etot_after - Etot_before + dLmag = norm2(Ltot_after(:) - Ltot_before(:)) + else + Ltot_before(:) = tmpsys%Lorbit(:) + tmpsys%Lspin(:) + Lmag_before = norm2(Ltot_before(:)) + ke_orbit_before = tmpsys%ke_orbit + ke_spin_before = tmpsys%ke_spin + pe_before = tmpsys%pe + Etot_before = ke_orbit_before + ke_spin_before + pe_before + end if + end associate + return + end subroutine calculate_system_energy + + subroutine shift_vector_to_origin(m_frag, vec_frag) + !! Author: Jennifer L.L. Pouplin, Carlisle A. Wishard, and David A. Minton + !! + !! Adjusts the position or velocity of the fragments as needed to align them with the origin + implicit none + ! Arguments + real(DP), dimension(:), intent(in) :: m_frag !! Fragment masses + real(DP), dimension(:,:), intent(inout) :: vec_frag !! Fragment positions or velocities in the center of mass frame + + ! Internals + real(DP), dimension(NDIM) :: mvec_frag, COM_offset + integer(I4B) :: i + + mvec_frag(:) = 0.0_DP + + do i = 1, nfrag + mvec_frag = mvec_frag(:) + vec_frag(:,i) * m_frag(i) + end do + COM_offset(:) = -mvec_frag(:) / mtot + do i = 1, nfrag + vec_frag(:, i) = vec_frag(:, i) + COM_offset(:) + end do + + return + end subroutine shift_vector_to_origin + + subroutine set_fragment_position_vectors() + !! Author: Jennifer L.L. Pouplin, Carlisle A. Wishard, and David A. Minton + !! + !! Initializes the orbits of the fragments around the center of mass. The fragments are initially placed on a plane defined by the + !! pre-impact angular momentum. They are distributed on an ellipse surrounding the center of mass. + !! The initial positions do not conserve energy or momentum, so these need to be adjusted later. + + implicit none + real(DP) :: dis, rad + real(DP), dimension(NDIM) :: L_sigma + logical, dimension(:), allocatable :: loverlap + integer(I4B) :: i, j + + allocate(loverlap(nfrag)) + + ! Place the fragments into a region that is big enough that we should usually not have overlapping bodies + ! An overlapping bodies will collide in the next time step, so it's not a major problem if they do (it just slows the run down) + r_max = r_max_start + rad = sum(radius(:)) + + ! We will treat the first two fragments of the list as special cases. They get initialized the maximum distances apart along the original impactor distance vector. + ! This is done because in a regular disruption, the first body is the largest, the second the second largest, and the rest are smaller equal-mass fragments. + + call random_number(x_frag(:,3:nfrag)) + loverlap(:) = .true. + do while (any(loverlap(3:nfrag))) + x_frag(:, 1) = x(:, 1) - xcom(:) + x_frag(:, 2) = x(:, 2) - xcom(:) + r_max = r_max + 0.1_DP * rad + do i = 3, nfrag + if (loverlap(i)) then + call random_number(x_frag(:,i)) + x_frag(:, i) = 2 * (x_frag(:, i) - 0.5_DP) * r_max + end if + end do + loverlap(:) = .false. + do j = 1, nfrag + do i = j + 1, nfrag + dis = norm2(x_frag(:,j) - x_frag(:,i)) + loverlap(i) = loverlap(i) .or. (dis <= (rad_frag(i) + rad_frag(j))) + end do + end do + end do + call shift_vector_to_origin(m_frag, x_frag) + + do i = 1, nfrag + rmag(i) = norm2(x_frag(:, i)) + v_r_unit(:, i) = x_frag(:, i) / rmag(i) + call random_number(L_sigma(:)) ! Randomize the tangential velocity direction. This helps to ensure that the tangential velocity doesn't completely line up with the angular momentum vector, + ! otherwise we can get an ill-conditioned system + v_h_unit(:, i) = z_col_unit(:) + 2e-1_DP * (L_sigma(:) - 0.5_DP) + v_h_unit(:, i) = v_h_unit(:, i) / norm2(v_h_unit(:, i)) + v_t_unit(:, i) = v_h_unit(:, i) .cross. v_r_unit(:, i) + xb_frag(:,i) = x_frag(:,i) + xcom(:) + end do + + return + end subroutine set_fragment_position_vectors + + subroutine set_fragment_tan_vel(lerr) + !! Author: Jennifer L.L. Pouplin, Carlisle A. Wishard, and David A. Minton + !! + !! Adjusts the tangential velocities and spins of a collection of fragments such that they conserve angular momentum without blowing the fragment kinetic energy budget. + !! This procedure works in several stages, with a goal to solve the angular and linear momentum constraints on the fragments, while still leaving a positive balance of + !! our fragment kinetic energy (ke_frag_budget) that we can put into the radial velocity distribution. + !! + !! The first thing we'll try to do is solve for the tangential velocities of the first 6 fragments, using angular and linear momentum as constraints and an initial + !! tangential velocity distribution for the remaining bodies (if there are any) that distributes their angular momentum equally between them. + !! If that doesn't work and we blow our kinetic energy budget, we will attempt to find a tangential velocity distribution that minimizes the kinetic energy while + !! conserving momentum. + !! + !! A failure will trigger a restructuring of the fragments so we will try new values of the radial position distribution. + implicit none + ! Arguments + logical, intent(out) :: lerr + ! Internals + integer(I4B) :: i + real(DP), parameter :: TOL = 1e-4_DP + real(DP), dimension(:), allocatable :: v_t_initial + type(lambda_obj) :: spinfunc + type(lambda_obj_err) :: objective_function + real(DP), dimension(NDIM) :: L_frag_spin, L_remainder, Li, rot_L, rot_ke + + ! Initialize the fragments with 0 velocity and spin so we can divide up the balance between the tangential, radial, and spin components while conserving momentum + lerr = .false. + vb_frag(:,:) = 0.0_DP + rot_frag(:,:) = 0.0_DP + v_t_mag(:) = 0.0_DP + v_r_mag(:) = 0.0_DP + + call calculate_system_energy(linclude_fragments=.true.) + ke_frag_budget = -dEtot - Qloss + !write(*,*) '***************************************************' + !write(*,*) 'Original dis : ',norm2(x(:,2) - x(:,1)) + !write(*,*) 'r_max : ',r_max + !write(*,*) 'f_spin : ',f_spin + !write(*,*) '***************************************************' + !write(*,*) 'Energy balance so far: ' + !write(*,*) 'ke_frag_budget : ',ke_frag_budget + !write(*,*) 'ke_orbit_before: ',ke_orbit_before + !write(*,*) 'ke_orbit_after : ',ke_orbit_after + !write(*,*) 'ke_spin_before : ',ke_spin_before + !write(*,*) 'ke_spin_after : ',ke_spin_after + !write(*,*) 'pe_before : ',pe_before + !write(*,*) 'pe_after : ',pe_after + !write(*,*) 'Qloss : ',Qloss + !write(*,*) '***************************************************' + if (ke_frag_budget < 0.0_DP) then + write(*,*) 'Negative ke_frag_budget: ',ke_frag_budget + r_max_start = r_max_start / 2 + lerr = .true. + return + end if + + allocate(v_t_initial, mold=v_t_mag) + + L_frag_spin(:) = 0.0_DP + ke_frag_spin = 0.0_DP + ! Start the first two bodies with the same rotation as the original two impactors, then distribute the remaining angular momentum among the rest + do i = 1, 2 + rot_frag(:, i) = rot(:, i) + L_frag_spin(:) = L_frag_spin(:) + m_frag(i) * rad_frag(i)**2 * Ip_frag(3, i) * rot_frag(:, i) + end do + L_frag_orb(:) = L_frag_tot(:) - L_frag_spin(:) + L_frag_spin(:) = 0.0_DP + do i = 1, nfrag + ! Convert a fraction (f_spin) of either the remaining angular momentum or kinetic energy budget into spin, whichever gives the smaller rotation so as not to blow any budgets + rot_ke(:) = sqrt(2 * f_spin * ke_frag_budget / (nfrag * m_frag(i) * rad_frag(i)**2 * Ip_frag(3, i))) * L_frag_orb(:) / norm2(L_frag_orb(:)) + rot_L(:) = f_spin * L_frag_orb(:) / (nfrag * m_frag(i) * rad_frag(i)**2 * Ip_frag(3, i)) + if (norm2(rot_ke) < norm2(rot_L)) then + rot_frag(:,i) = rot_frag(:, i) + rot_ke(:) + else + rot_frag(:, i) = rot_frag(:, i) + rot_L(:) + end if + L_frag_spin(:) = L_frag_spin(:) + m_frag(i) * rad_frag(i)**2 * Ip_frag(3, i) * rot_frag(:, i) + ke_frag_spin = ke_frag_spin + m_frag(i) * Ip_frag(3, i) * rad_frag(i)**2 * dot_product(rot_frag(:, i), rot_frag(:, i)) + end do + ke_frag_spin = 0.5_DP * ke_frag_spin + ! Convert a fraction of the pre-impact angular momentum into fragment spin angular momentum + L_frag_orb(:) = L_frag_tot(:) - L_frag_spin(:) + L_remainder(:) = L_frag_orb(:) + ! Next we will solve for the tangential component of the velocities that both conserves linear momentum and uses the remaining angular momentum not used in spin. + ! This will be done using a linear solver that solves for the tangential velocities of the first 6 fragments, constrained by the linear and angular momentum vectors, + ! which is embedded in a non-linear minimizer that will adjust the tangential velocities of the remaining i>6 fragments to minimize kinetic energy for a given momentum solution + ! The initial conditions fed to the minimizer for the fragments will be the remaining angular momentum distributed between the fragments. + do i = 1, nfrag + v_t_initial(i) = norm2(L_remainder(:)) / ((nfrag - i + 1) * m_frag(i) * norm2(x_frag(:,i))) + Li(:) = m_frag(i) * x_frag(:,i) .cross. v_t_initial(i) * v_t_unit(:, i) + L_remainder(:) = L_remainder(:) - Li(:) + end do + + ! Find the local kinetic energy minimum for the system that conserves linear and angular momentum + objective_function = lambda_obj(tangential_objective_function, lerr) + v_t_mag(7:nfrag) = util_minimize_bfgs(objective_function, nfrag-6, v_t_initial(7:nfrag), TOL, lerr) + ! Now that the KE-minimized values of the i>6 fragments are found, calculate the momentum-conserving solution for tangential velociteis + v_t_initial(7:nfrag) = v_t_mag(7:nfrag) + v_t_mag(1:nfrag) = solve_fragment_tan_vel(v_t_mag_input=v_t_initial(7:nfrag), lerr=lerr) + + ! Perform one final shift of the radial velocity vectors to align with the center of mass of the collisional system (the origin) + vb_frag(:,1:nfrag) = vmag_to_vb(v_r_mag(1:nfrag), v_r_unit(:,1:nfrag), v_t_mag(1:nfrag), v_t_unit(:,1:nfrag), m_frag(1:nfrag), vcom(:)) + ! Now do a kinetic energy budget check to make sure we are still within the budget. + ke_frag_orbit = 0.0_DP + do i = 1, nfrag + v_frag(:, i) = vb_frag(:, i) - vcom(:) + ke_frag_orbit = ke_frag_orbit + m_frag(i) * dot_product(vb_frag(:, i), vb_frag(:, i)) + end do + ke_frag_orbit = 0.5_DP * ke_frag_orbit + ke_radial = ke_frag_budget - ke_frag_orbit - ke_frag_spin + + ! If we are over the energy budget, flag this as a failure so we can try again + lerr = (ke_radial < 0.0_DP) + !write(*,*) 'Tangential' + !write(*,*) 'ke_frag_budget: ',ke_frag_budget + !write(*,*) 'ke_frag_orbit : ',ke_frag_orbit + !write(*,*) 'ke_frag_spin : ',ke_frag_spin + !write(*,*) 'ke_radial : ',ke_radial + + return + + end subroutine set_fragment_tan_vel + + function tangential_objective_function(v_t_mag_input, lerr) result(fval) + !! Author: David A. Minton + !! + !! 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_t_mag_input !! Unknown tangential component of velocity vector set previously by angular momentum constraint + logical, intent(out) :: lerr !! Error flag + ! Result + real(DP) :: fval + ! Internals + integer(I4B) :: i + real(DP), dimension(:,:), allocatable :: v_shift + real(DP), dimension(:), allocatable :: v_t_new + real(DP) :: keo + + lerr = .false. + + allocate(v_shift(NDIM, nfrag)) + allocate(v_t_new(nfrag)) + + v_t_new(:) = solve_fragment_tan_vel(v_t_mag_input=v_t_mag_input(:), lerr=lerr) + v_shift(:,:) = vmag_to_vb(v_r_mag, v_r_unit, v_t_new, v_t_unit, m_frag, vcom) + + keo = 0.0_DP + do i = 1, nfrag + keo = keo + m_frag(i) * dot_product(v_shift(:, i), v_shift(:, i)) + end do + keo = 0.5_DP * keo + fval = keo + lerr = .false. + return + end function tangential_objective_function + + function solve_fragment_tan_vel(lerr, v_t_mag_input) result(v_t_mag_output) + !! Author: Jennifer L.L. Pouplin, Carlisle A. Wishard, and David A. Minton + !! + !! Adjusts the positions, velocities, and spins of a collection of fragments such that they conserve angular momentum + implicit none + ! Arguments + logical, intent(out) :: lerr !! Error flag + real(DP), dimension(:), optional, intent(in) :: v_t_mag_input !! Unknown tangential velocities for fragments 7:nfrag + ! Internals + integer(I4B) :: i + ! Result + real(DP), dimension(:), allocatable :: v_t_mag_output + + real(DP), dimension(2 * NDIM, 2 * NDIM) :: A ! LHS of linear equation used to solve for momentum constraint in Gauss elimination code + real(DP), dimension(2 * NDIM) :: b ! RHS of linear equation used to solve for momentum constraint in Gauss elimination code + real(DP), dimension(NDIM) :: L_lin_others, L_orb_others, L, vtmp + + v_frag(:,:) = 0.0_DP + lerr = .false. + + ! We have 6 constraint equations (2 vector constraints in 3 dimensions each) + ! The first 3 are that the linear momentum of the fragments is zero with respect to the collisional barycenter + ! The second 3 are that the sum of the angular momentum of the fragments is conserved from the pre-impact state + L_lin_others(:) = 0.0_DP + L_orb_others(:) = 0.0_DP + do i = 1, nfrag + if (i <= 2 * NDIM) then ! The tangential velocities of the first set of bodies will be the unknowns we will solve for to satisfy the constraints + A(1:3, i) = m_frag(i) * v_t_unit(:, i) + L(:) = v_r_unit(:, i) .cross. v_t_unit(:, i) + A(4:6, i) = m_frag(i) * rmag(i) * L(:) + else if (present(v_t_mag_input)) then + vtmp(:) = v_t_mag_input(i - 6) * v_t_unit(:, i) + L_lin_others(:) = L_lin_others(:) + m_frag(i) * vtmp(:) + L(:) = x_frag(:, i) .cross. vtmp(:) + L_orb_others(:) = L_orb_others(:) + m_frag(i) * L(:) + end if + end do + b(1:3) = -L_lin_others(:) + b(4:6) = L_frag_orb(:) - L_orb_others(:) + allocate(v_t_mag_output(nfrag)) + v_t_mag_output(1:6) = util_solve_linear_system(A, b, 6, lerr) + if (present(v_t_mag_input)) v_t_mag_output(7:nfrag) = v_t_mag_input(:) + + return + end function solve_fragment_tan_vel + + subroutine set_fragment_radial_velocities(lerr) + !! Author: Jennifer L.L. Pouplin, Carlisle A. Wishard, and David A. Minton + !! + !! + !! Adjust the fragment velocities to set the fragment orbital kinetic energy. This will minimize the difference between the fragment kinetic energy and the energy budget + implicit none + ! Arguments + logical, intent(out) :: lerr + ! Internals + real(DP), parameter :: TOL = 1e-10_DP + integer(I4B) :: i, j + real(DP), dimension(:), allocatable :: v_r_initial, v_r_sigma + real(DP), dimension(:,:), allocatable :: v_r + type(lambda_obj) :: objective_function + + ! Set the "target" ke_orbit_after (the value of the orbital kinetic energy that the fragments ought to have) + + allocate(v_r_initial, source=v_r_mag) + ! Initialize radial velocity magnitudes with a random value that is approximately 10% of that found by distributing the kinetic energy equally + allocate(v_r_sigma, source=v_r_mag) + call random_number(v_r_sigma(1:nfrag)) + v_r_sigma(1:nfrag) = sqrt(1.0_DP + 2 * (v_r_sigma(1:nfrag) - 0.5_DP) * 1e-4_DP) + v_r_initial(1:nfrag) = v_r_sigma(1:nfrag) * sqrt(abs(ke_radial) / (2 * m_frag(1:nfrag) * nfrag)) + + ! Initialize the lambda function using a structure constructor that calls the init method + ! Minimize the ke objective function using the BFGS optimizer + objective_function = lambda_obj(radial_objective_function) + v_r_mag = util_minimize_bfgs(objective_function, nfrag, v_r_initial, TOL, lerr) + ! Shift the radial velocity vectors to align with the center of mass of the collisional system (the origin) + vb_frag(:,1:nfrag) = vmag_to_vb(v_r_mag(1:nfrag), v_r_unit(:,1:nfrag), v_t_mag(1:nfrag), v_t_unit(:,1:nfrag), m_frag(1:nfrag), vcom(:)) + do i = 1, nfrag + v_frag(:, i) = vb_frag(:, i) - vcom(:) + end do + ke_frag_orbit = 0.0_DP + do i = 1, nfrag + ke_frag_orbit = ke_frag_orbit + m_frag(i) * dot_product(vb_frag(:, i), vb_frag(:, i)) + end do + ke_frag_orbit = 0.5_DP * ke_frag_orbit + !write(*,*) 'Radial' + !write(*,*) 'Failure? ',lerr + !write(*,*) 'ke_frag_budget: ',ke_frag_budget + !write(*,*) 'ke_frag_orbit : ',ke_frag_orbit + !write(*,*) 'ke_frag_spin : ',ke_frag_spin + !write(*,*) 'ke_remainder : ',ke_frag_budget - (ke_frag_orbit + ke_frag_spin) + lerr = .false. + + return + end subroutine set_fragment_radial_velocities + + function radial_objective_function(v_r_mag_input) result(fval) + !! Author: David A. Minton + !! + !! 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_input !! Unknown radial component of fragment velocity vector + ! Result + real(DP) :: fval !! The objective function result, which is the square of the difference between the calculated fragment kinetic energy and our target + !! Minimizing this brings us closer to our objective + ! Internals + integer(I4B) :: i + real(DP), dimension(:,:), allocatable :: v_shift + + allocate(v_shift, mold=vb_frag) + v_shift(:,:) = vmag_to_vb(v_r_mag_input, v_r_unit, v_t_mag, v_t_unit, m_frag, vcom) + fval = 2 * ke_frag_budget + do i = 1, nfrag + fval = fval - m_frag(i) * (Ip_frag(3, i) * rad_frag(i)**2 * dot_product(rot_frag(:, i), rot_frag(:, i)) + dot_product(v_shift(:, i), v_shift(:, i))) + end do + ! The following ensures that fval = 0 is a local minimum, which is what the BFGS method is searching for + fval = (fval / (2 * ke_radial))**2 + + return + + end function radial_objective_function + + function vmag_to_vb(v_r_mag, v_r_unit, v_t_mag, v_t_unit, m_frag, vcom) result(vb) + !! Author: David A. Minton + !! + !! Converts radial and tangential velocity magnitudes into barycentric velocity + implicit none + ! Arguments + real(DP), dimension(:), intent(in) :: v_r_mag !! Unknown radial component of fragment velocity vector + real(DP), dimension(:), intent(in) :: v_t_mag !! Tangential component of velocity vector set previously by angular momentum constraint + real(DP), dimension(:,:), intent(in) :: v_r_unit, v_t_unit !! Radial and tangential unit vectors for each fragment + real(DP), dimension(:), intent(in) :: m_frag !! Fragment masses + real(DP), dimension(:), intent(in) :: vcom !! Barycentric velocity of collisional system center of mass + ! Result + real(DP), dimension(:,:), allocatable :: vb + ! Internals + integer(I4B) :: i + + allocate(vb, mold=v_r_unit) + ! Make sure the velocity magnitude stays positive + do i = 1, nfrag + vb(:,i) = abs(v_r_mag(i)) * v_r_unit(:, i) + end do + ! In order to keep satisfying the kinetic energy constraint, we must shift the origin of the radial component of the velocities to the center of mass + call shift_vector_to_origin(m_frag, vb) + + do i = 1, nfrag + vb(:, i) = vb(:, i) + v_t_mag(i) * v_t_unit(:, i) + vcom(:) + end do + + end function vmag_to_vb + + subroutine restructure_failed_fragments() + !! Author: David A. Minton + !! + !! We failed to find a set of positions and velocities that satisfy all the constraints, and so we will alter the fragments and try again. + implicit none + integer(I4B) :: i + real(DP), dimension(:), allocatable :: m_frag_new, rad_frag_new + real(DP), dimension(:,:), allocatable :: xb_frag_new, vb_frag_new, Ip_frag_new, rot_frag_new + real(DP) :: delta_r, delta_r_max + real(DP), parameter :: ke_avg_deficit_target = 0.0_DP + + ! Introduce a bit of noise in the radius determination so we don't just flip flop between similar failed positions + call random_number(delta_r_max) + delta_r_max = sum(radius(:)) * (1.0_DP + 2e-1_DP * (delta_r_max - 0.5_DP)) + if (try > 2) then + ! Linearly interpolate the last two failed solution ke deficits to find a new distance value to try + delta_r = (r_max_start - r_max_start_old) * (ke_avg_deficit_target - ke_avg_deficit_old) / (ke_avg_deficit - ke_avg_deficit_old) + if (abs(delta_r) > delta_r_max) delta_r = sign(delta_r_max, delta_r) + else + delta_r = delta_r_max + end if + r_max_start_old = r_max_start + r_max_start = r_max_start + delta_r ! The larger lever arm can help if the problem is in the angular momentum step + if (f_spin > epsilon(1.0_DP)) then + f_spin = f_spin / 2 + else + f_spin = 0.0_DP + end if + end subroutine restructure_failed_fragments + + + end subroutine fragmentation_initialize + + + module subroutine fragmentation_regime(Mcb, m1, m2, rad1, rad2, xh1, xh2, vb1, vb2, den1, den2, regime, Mlr, Mslr, mtiny, Qloss) !! Author: Jennifer L.L. Pouplin, Carlisle A. Wishard, and David A. Minton !! diff --git a/src/modules/swiftest_classes.f90 b/src/modules/swiftest_classes.f90 index 1aec7dc9d..d00d271fb 100644 --- a/src/modules/swiftest_classes.f90 +++ b/src/modules/swiftest_classes.f90 @@ -966,6 +966,17 @@ end subroutine util_fill_arr_logical end interface interface + module function util_minimize_bfgs(f, N, x0, eps, lerr) result(x1) + use lambda_function + implicit none + integer(I4B), intent(in) :: N + class(lambda_obj), intent(inout) :: 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 + module subroutine util_peri_tp(self, system, param) implicit none class(swiftest_tp), intent(inout) :: self !! Swiftest test particle object @@ -974,6 +985,39 @@ module subroutine util_peri_tp(self, system, param) end subroutine util_peri_tp end interface + interface util_solve_linear_system + module function util_solve_linear_system_d(A,b,n,lerr) result(x) + implicit none + integer(I4B), intent(in) :: n + real(DP), dimension(:,:), intent(in) :: A + real(DP), dimension(:), intent(in) :: b + logical, intent(out) :: lerr + real(DP), dimension(n) :: x + end function util_solve_linear_system_d + + module function util_solve_linear_system_q(A,b,n,lerr) result(x) + implicit none + integer(I4B), intent(in) :: n + real(QP), dimension(:,:), intent(in) :: A + real(QP), dimension(:), intent(in) :: b + logical, intent(out) :: lerr + real(QP), dimension(n) :: x + end function util_solve_linear_system_q + end interface + + interface + module function util_solve_rkf45(f, y0in, t1, dt0, tol) result(y1) + use lambda_function + implicit none + class(lambda_obj), intent(inout) :: f !! lambda function object that has been initialized to be a function of derivatives. The object will return with components lastarg and lasteval set + real(DP), dimension(:), intent(in) :: y0in !! Initial value at t=0 + real(DP), intent(in) :: t1 !! Final time + real(DP), intent(in) :: dt0 !! Initial step size guess + real(DP), intent(in) :: tol !! Tolerance on solution + real(DP), dimension(:), allocatable :: y1 !! Final result + end function util_solve_rkf45 + end interface + interface util_resize module subroutine util_resize_arr_char_string(arr, nnew) implicit none diff --git a/src/util/util_minimize_bfgs.f90 b/src/util/util_minimize_bfgs.f90 new file mode 100644 index 000000000..fbd48c8c2 --- /dev/null +++ b/src/util/util_minimize_bfgs.f90 @@ -0,0 +1,584 @@ +submodule (swiftest_classes) s_util_minimize_bfgs + use swiftest +contains + module function util_minimize_bfgs(f, N, x0, eps, lerr) result(x1) + !! author: David A. Minton + !! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - + !! This function implements the Broyden-Fletcher-Goldfarb-Shanno method to determine the minimum of a function of N variables. + !! It recieves as input: + !! f%eval(x) : lambda function object containing the objective function as the eval metho + !! N : Number of variables of function f + !! x0 : Initial starting value of x + !! eps : Accuracy of 1 - dimensional minimization at each step + !! The outputs include + !! lerr : Returns .true. if it could not find the minimum + !! Returns + !! x1 : Final minimum (all 0 if none found) + !! 0 = No miniumum found + !! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - + use, intrinsic :: ieee_exceptions + implicit none + ! Arguments + integer(I4B), intent(in) :: N + class(lambda_obj), intent(inout) :: f + real(DP), dimension(:), intent(in) :: x0 + real(DP), intent(in) :: eps + logical, intent(out) :: lerr + ! Result + real(DP), dimension(:), allocatable :: x1 + ! Internals + integer(I4B) :: i, j, k, l, conv, num + integer(I4B), parameter :: MAXLOOP = 1000 !! Maximum number of loops before method is determined to have failed + real(DP), parameter :: graddelta = 1e-4_DP !! Delta x for gradient calculations + real(DP), dimension(N) :: S !! Direction vectors + real(DP), dimension(N,N) :: H !! Approximated inverse Hessian matrix + real(DP), dimension(N) :: grad1 !! gradient of f + real(DP), dimension(N) :: grad0 !! old value of gradient + real(DP) :: astar !! 1D minimized value + real(DP), dimension(N) :: y, P + real(DP), dimension(N,N) :: PP, PyH, HyP + real(DP) :: yHy, Py + type(ieee_status_type) :: original_fpe_status + logical, dimension(:), allocatable :: fpe_flag + + call ieee_get_status(original_fpe_status) ! Save the original floating point exception status + call ieee_set_flag(ieee_all, .false.) ! Set all flags to quiet + allocate(fpe_flag(size(ieee_usual))) + + lerr = .false. + allocate(x1, source=x0) + ! Initialize approximate Hessian with the identity matrix (i.e. begin with method of steepest descent) + ! Get initial gradient and initialize arrays for updated values of gradient and x + H(:,:) = reshape([((0._DP, i=1, j-1), 1._DP, (0._DP, i=j+1, N), j=1, N)], [N,N]) + grad0 = gradf(f, N, x0(:), graddelta, lerr) + if (lerr) then + call ieee_set_status(original_fpe_status) + return + end if + grad1(:) = grad0(:) + do i = 1, MAXLOOP + !check for convergence + conv = count(abs(grad1(:)) > eps) + if (conv == 0) then + !write(*,*) "BFGS converged on gradient after ",i," iterations" + exit + end if + S(:) = -matmul(H(:,:), grad1(:)) + astar = minimize1D(f, x1, S, N, graddelta, lerr) + if (lerr) then + !write(*,*) "Exiting BFGS with error in minimize1D step" + exit + end if + ! Get new x values + P(:) = astar * S(:) + x1(:) = x1(:) + P(:) + ! Calculate new gradient + grad0(:) = grad1(:) + grad1 = gradf(f, N, x1, graddelta, lerr) + y(:) = grad1(:) - grad0(:) + Py = sum(P(:) * y(:)) + ! set up factors for H matrix update + yHy = 0._DP + do k = 1, N + do j = 1, N + yHy = yHy + y(j) * H(j,k) * y(k) + end do + end do + ! prevent divide by zero (convergence) + if (abs(Py) < tiny(Py)) then + !write(*,*) "BFGS Converged on tiny Py after ",i," iterations" + exit + end if + ! set up update + PyH(:,:) = 0._DP + HyP(:,:) = 0._DP + do k = 1, N + do j = 1, N + PP(j, k) = P(j) * P(k) + do l = 1, N + PyH(j, k) = PyH(j, k) + P(j) * y(l) * H(l,k) + HyP(j, k) = HyP(j, k) + P(k) * y(l) * H(j,l) + end do + end do + end do + ! update H matrix + H(:,:) = H(:,:) + ((1._DP - yHy / Py) * PP(:,:) - PyH(:,:) - HyP(:,:)) / Py + ! Normalize to prevent it from blowing up if it takes many iterations to find a solution + H(:,:) = H(:,:) / norm2(H(:,:)) + ! Stop everything if there are any exceptions to allow the routine to fail gracefully + call ieee_get_flag(ieee_usual, fpe_flag) + if (any(fpe_flag)) exit + if (i == MAXLOOP) then + lerr = .true. + !write(*,*) "BFGS ran out of loops!" + end if + end do + call ieee_get_flag(ieee_usual, fpe_flag) + lerr = lerr .or. any(fpe_flag) + !if (any(fpe_flag)) write(*,*) 'BFGS did not converge due to fpe' + !if (lerr) write(*,*) "BFGS did not converge!" + call ieee_set_status(original_fpe_status) + + return + + contains + + function gradf(f, N, x1, dx, lerr) result(grad) + !! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - + !! Purpose: Estimates the gradient of a function using a central difference + !! approximation + !! Inputs: + !! f%eval(x) : lambda function object containing the objective function as the eval metho + !! N : number of variables N + !! x1 : x value array + !! dx : step size to use when calculating derivatives + !! Outputs: + !! lerr : .true. if an error occurred. Otherwise returns .false. + !! Returns + !! grad : N sized array containing estimated gradient of f at x1 + !! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - + implicit none + ! Arguments + integer(I4B), intent(in) :: N + class(lambda_obj), intent(inout) :: f + real(DP), dimension(:), intent(in) :: x1 + real(DP), intent(in) :: dx + logical, intent(out) :: lerr + ! Result + real(DP), dimension(N) :: grad + ! Internals + integer(I4B) :: i, j + real(DP), dimension(N) :: xp, xm + real(DP) :: fp, fm + logical :: lerrp, lerrm + + do i = 1, N + do j = 1, N + if (j == i) then + xp(j) = x1(j) + dx + xm(j) = x1(j) - dx + else + xp(j) = x1(j) + xm(j) = x1(j) + end if + end do + select type (f) + class is (lambda_obj_err) + fp = f%eval(xp) + lerrp = f%lerr + fm = f%eval(xm) + lerrm = f%lerr + lerr = lerrp .or. lerrm + class is (lambda_obj) + fp = f%eval(xp) + fm = f%eval(xm) + lerr = .false. + end select + grad(i) = (fp - fm) / (2 * dx) + if (lerr) return + end do + return + end function gradf + + function minimize1D(f, x0, S, N, eps, lerr) result(astar) + !! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - + !! This program find the minimum of a function of N variables in a single direction + !! S using in sequence: + !! 1. A Bracketing method + !! 2. The golden section method + !! 3. A quadratic polynomial fit + !! Inputs + !! f%eval(x) : lambda function object containing the objective function as the eval metho + !! x0 : Array of size N of initial x values + !! S : Array of size N that determines the direction of minimization + !! N : Number of variables of function f + !! eps : Accuracy of 1 - dimensional minimization at each step + !! Output + !! lerr : .true. if an error occurred. Otherwise returns .false. + !! Returns + !! astar : Final minimum along direction S + !! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - + implicit none + ! Arguments + integer(I4B), intent(in) :: N + class(lambda_obj), intent(inout) :: f + real(DP), dimension(:), intent(in) :: x0, S + real(DP), intent(in) :: eps + logical, intent(out) :: lerr + ! Result + real(DP) :: astar + ! Internals + integer(I4B) :: num = 0 + real(DP), parameter :: step = 0.7_DP !! Bracketing method step size + real(DP), parameter :: gam = 1.2_DP !! Bracketing method expansion parameter + real(DP), parameter :: greduce = 0.2_DP !! Golden section method reduction factor + real(DP), parameter :: greduce2 = 0.1_DP ! Secondary golden section method reduction factor + real(DP) :: alo, ahi !! High and low values for 1 - D minimization routines + real(DP), parameter :: a0 = epsilon(1.0_DP) !! Initial guess of alpha + + alo = a0 + call bracket(f, x0, S, N, gam, step, alo, ahi, lerr) + if (lerr) then + !write(*,*) "BFGS bracketing step failed!" + return + end if + if (abs(alo - ahi) < eps) then + astar = alo + lerr = .false. + return + end if + call golden(f, x0, S, N, greduce, alo, ahi, lerr) + if (lerr) then + !write(*,*) "BFGS golden section step failed!" + return + end if + if (abs(alo - ahi) < eps) then + astar = alo + lerr = .false. + return + end if + call quadfit(f, x0, S, N, eps, alo, ahi, lerr) + if (lerr) then + !write(*,*) "BFGS quadfit failed!" + return + end if + if (abs(alo - ahi) < eps) then + astar = alo + lerr = .false. + return + end if + ! Quadratic fit method won't converge, so finish off with another golden section + call golden(f, x0, S, N, greduce2, alo, ahi, lerr) + if (.not. lerr) astar = (alo + ahi) / 2.0_DP + return + end function minimize1D + + function n2one(f, x0, S, N, a, lerr) result(fnew) + implicit none + ! Arguments + integer(I4B), intent(in) :: N + class(lambda_obj), intent(inout) :: f + real(DP), dimension(:), intent(in) :: x0, S + real(DP), intent(in) :: a + logical, intent(out) :: lerr + + ! Return + real(DP) :: fnew + ! Internals + real(DP), dimension(N) :: xnew + integer(I4B) :: i + + xnew(:) = x0(:) + a * S(:) + fnew = f%eval(xnew(:)) + select type(f) + class is (lambda_obj_err) + lerr = f%lerr + class is (lambda_obj) + lerr = .false. + end select + return + end function n2one + + ! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - + ! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - + + subroutine bracket(f, x0, S, N, gam, step, lo, hi, lerr) + ! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - + !! This subroutine brackets the minimum. It recieves as input: + !! f%eval(x) : lambda function object containing the objective function as the eval metho + !! x0 : Array of size N of initial x values + !! S : Array of size N that determines the direction of minimization + !! gam : expansion parameter + !! step : step size + !! lo : initial guess of lo bracket value + !! The outputs include + !! lo : lo bracket + !! hi : hi bracket + !! lerr : .true. if an error occurred. Otherwise returns .false. + !! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - + implicit none + ! Arguments + integer(I4B), intent(in) :: N + class(lambda_obj), intent(inout) :: f + real(DP), dimension(:), intent(in) :: x0, S + real(DP), intent(in) :: gam, step + real(DP), intent(inout) :: lo + real(DP), intent(out) :: hi + logical, intent(out) :: lerr + ! Internals + real(DP) :: a0, a1, a2, atmp, da + real(DP) :: f0, f1, f2 + integer(I4B) :: i, j + integer(I4B), parameter :: MAXLOOP = 100 ! maximum number of loops before method is determined to have failed + real(DP), parameter :: eps = epsilon(lo) ! small number precision to test floating point equality + + ! set up initial bracket points + a0 = lo + da = step + a1 = a0 + da + a2 = a0 + 2 * da + f0 = n2one(f, x0, S, N, a0, lerr) + if (lerr) return + f1 = n2one(f, x0, S, N, a1, lerr) + if (lerr) return + f2 = n2one(f, x0, S, N, a2, lerr) + if (lerr) return + ! loop over bracket method until either min is bracketed method fails + do i = 1, MAXLOOP + if ((f0 > f1) .and. (f1 < f2)) then ! Minimum was found + lo = a0 + hi = a2 + return + else if ((f0 >= f1) .and. (f1 > f2)) then ! Function appears to decrease + da = da * gam + atmp = a2 + da + a0 = a1 + a1 = a2 + a2 = atmp + f0 = f1 + f1 = f2 + f2 = n2one(f, x0, S, N, a2, lerr) + else if ((f0 < f1) .and. (f1 <= f2)) then ! Function appears to increase + da = da * gam + atmp = a0 - da + a2 = a1 + a1 = a0 + a0 = atmp + f2 = f1 + f0 = n2one(f, x0, S, N, a0, lerr) + else if ((f0 < f1) .and. (f1 > f2)) then ! We are at a peak. Pick the direction that descends the fastest + da = da * gam + if (f2 > f0) then ! LHS is lower than RHS + atmp = a2 + da + a0 = a1 + a1 = a2 + a2 = atmp + f0 = f1 + f1 = f2 + f2 = n2one(f, x0, S, N, a2, lerr) + else ! RHS is lower than LHS + atmp = a0 - da + a2 = a1 + a1 = a0 + a0 = atmp + f2 = f1 + f1 = f2 + f0 = n2one(f, x0, S, N, a0, lerr) + end if + else if ((f0 > f1) .and. (abs(f2 - f1) <= eps)) then ! Decrasging but RHS equal + da = da * gam + atmp = a2 + da + a2 = atmp + f2 = n2one(f, x0, S, N, a2, lerr) + else if ((abs(f0 - f1) < eps) .and. (f1 < f2)) then ! Increasing but LHS equal + da = da * gam + atmp = a0 - da + a0 = atmp + f0 = n2one(f, x0, S, N, a0, lerr) + else ! all values equal. Expand in either direction and try again + a0 = a0 - da + a2 = a2 + da + f0 = n2one(f, x0, S, N, a0, lerr) + if (lerr) exit ! An error occurred while evaluating the function + f2 = n2one(f, x0, S, N, a2, lerr) + end if + if (lerr) exit ! An error occurred while evaluating the function + end do + lerr = .true. + return ! no minimum found + end subroutine bracket + + ! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - + ! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - + + subroutine golden(f, x0, S, N, eps, lo, hi, lerr) + ! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - + !! This function uses the golden section method to reduce the starting interval lo, hi by some amount sigma. + !! It recieves as input: + !! f%eval(x) : lambda function object containing the objective function as the eval metho + !! x0 : Array of size N of initial x values + !! S : Array of size N that determines the direction of minimization + !! gam : expansion parameter + !! eps : reduction interval in range (0 < sigma < 1) such that: + !! hi(new) - lo(new) = eps * (hi(old) - lo(old)) + !! lo : initial guess of lo bracket value + !! The outputs include + !! lo : lo bracket + !! hi : hi bracket + !! lerr : .true. if an error occurred. Otherwise returns .false. + !! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - + implicit none + ! Arguments + integer(I4B), intent(in) :: N + class(lambda_obj), intent(inout) :: f + real(DP), dimension(:), intent(in) :: x0, S + real(DP), intent(in) :: eps + real(DP), intent(inout) :: lo + real(DP), intent(out) :: hi + logical, intent(out) :: lerr + ! Internals + real(DP), parameter :: tau = 0.5_DP * (sqrt(5.0_DP) - 1.0_DP) ! Golden section constant + integer(I4B), parameter :: MAXLOOP = 40 ! maximum number of loops before method is determined to have failed (unlikely, but could occur if no minimum exists between lo and hi) + real(DP) :: i0 ! Initial interval value + real(DP) :: a1, a2 + real(DP) :: f1, f2 + integer(I4B) :: i, j + + i0 = hi - lo + a1 = hi - tau * i0 + a2 = lo + tau * i0 + f1 = n2one(f, x0, S, N, a1, lerr) + if (lerr) return + f2 = n2one(f, x0, S, N, a2, lerr) + if (lerr) return + do i = 1, MAXLOOP + if (abs((hi - lo) / i0) <= eps) return ! interval reduced to input amount + if (f2 > f1) then + hi = a2 + a2 = a1 + f2 = f1 + a1 = hi - tau * (hi - lo) + f1 = n2one(f, x0, S, N, a1, lerr) + else + lo = a1 + a1 = a2 + f2 = f1 + a2 = hi - (1.0_DP - tau) * (hi - lo) + f2 = n2one(f, x0, S, N, a2, lerr) + end if + if (lerr) exit + end do + lerr = .true. + return ! search took too many iterations - no minimum found + end subroutine golden + + ! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - + ! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - + + subroutine quadfit(f, x0, S, N, eps, lo, hi, lerr) + ! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - + !! This function uses a quadratic polynomial fit to locate the minimum of a function + !! to some accuracy eps. It recieves as input: + !! f%eval(x) : lambda function object containing the objective function as the eval metho + !! lo : low bracket value + !! hi : high bracket value + !! eps : desired accuracy of final minimum location + !! The outputs include + !! lo : final minimum location + !! hi : final minimum location + !! Notes: Uses the ieee_exceptions intrinsic module to allow for graceful failure due to floating point exceptions, which won't terminate the run. + !! - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - + implicit none + ! Arguments + integer(I4B), intent(in) :: N + class(lambda_obj), intent(inout) :: f + real(DP), dimension(:), intent(in) :: x0, S + real(DP), intent(in) :: eps + real(DP), intent(inout) :: lo + real(DP), intent(out) :: hi + logical, intent(out) :: lerr + ! Internals + integer(I4B), parameter :: MAXLOOP = 20 ! maximum number of loops before method is determined to have failed. + real(DP) :: a1, a2, a3, astar ! three points for the polynomial fit and polynomial minimum + real(DP) :: f1, f2, f3, fstar ! three function values for the polynomial and polynomial minimum + real(DP), dimension(3) :: row_1, row_2, row_3, rhs, soln ! matrix for 3 equation solver (gaussian elimination) + real(DP), dimension(3,3) :: lhs + real(DP) :: d1, d2, d3, aold, denom, errval + integer(I4B) :: i + + lerr = .false. + ! Get initial a1, a2, a3 values + a1 = lo + a2 = lo + 0.5_DP * (hi - lo) + a3 = hi + aold = a1 + astar = a2 + f1 = n2one(f, x0, S, N, a1, lerr) + if (lerr) return + f2 = n2one(f, x0, S, N, a2, lerr) + if (lerr) return + f3 = n2one(f, x0, S, N, a3, lerr) + if (lerr) return + do i = 1, MAXLOOP + ! check to see if convergence is reached and exit + errval = abs((astar - aold) / astar) + call ieee_get_flag(ieee_usual, fpe_flag) + if (any(fpe_flag)) then + !write(*,*) 'quadfit fpe' + !write(*,*) 'aold : ',aold + !write(*,*) 'astar: ',astar + lerr = .true. + exit + end if + if (errval < eps) then + lo = astar + hi = astar + exit + end if + ! Set up system for gaussian elimination equation solver + row_1 = [1.0_DP, a1, a1**2] + row_2 = [1.0_DP, a2, a2**2] + row_3 = [1.0_DP, a3, a3**2] + rhs = [f1, f2, f3] + lhs(1, :) = row_1 + lhs(2, :) = row_2 + lhs(3, :) = row_3 + ! Solve system of equations + soln(:) = util_solve_linear_system(lhs, rhs, 3, lerr) + call ieee_set_flag(ieee_all, .false.) ! Set all flags back to quiet + call ieee_set_halting_mode(ieee_divide_by_zero, .false.) + if (lerr) then + !write(*,*) 'quadfit fpe:' + !write(*,*) 'util_solve_linear_system failed' + exit + end if + aold = astar + if (soln(2) == soln(3)) then ! Handles the case where they are both 0. 0/0 is an unhandled exception + astar = -0.5_DP + else + astar = -soln(2) / (2 * soln(3)) + end if + call ieee_get_flag(ieee_usual, fpe_flag) + if (any(fpe_flag)) then + !write(*,*) 'quadfit fpe' + !write(*,*) 'soln(2:3): ',soln(2:3) + !write(*,*) 'a1, a2, a3' + !write(*,*) a1, a2, a3 + !write(*,*) 'f1, f2, f3' + !write(*,*) f1, f2, f3 + lerr = .true. + exit + end if + fstar = n2one(f, x0, S, N, astar, lerr) + if (lerr) exit + ! keep the three closest a values to astar and discard the fourth + d1 = abs(a1 - astar) + d2 = abs(a2 - astar) + d3 = abs(a3 - astar) + + if (d1 > d2) then + if (d1 > d3) then + f1 = fstar + a1 = astar + else if (d3 > d2) then + f3 = fstar + a3 = astar + end if + else + if (d2 > d3) then + f2 = fstar + a2 = astar + else if (d3 > d1) then + f3 = fstar + a3 = astar + end if + end if + end do + if (lerr) return + lo = a1 + hi = a3 + return + end subroutine quadfit + + end function util_minimize_bfgs +end submodule s_util_minimize_bfgs \ No newline at end of file diff --git a/src/util/util_solve.f90 b/src/util/util_solve.f90 index 0c3161ae2..ed9ff40ea 100644 --- a/src/util/util_solve.f90 +++ b/src/util/util_solve.f90 @@ -2,7 +2,143 @@ use swiftest contains - function util_solve_rkf45(f, y0in, t1, dt0, tol) result(y1) + module function util_solve_linear_system_d(A,b,n,lerr) result(x) + !! Author: David A. Minton + !! + !! Solves the linear equation of the form A*x = b for x. + !! A is an (n,n) arrays + !! x and b are (n) arrays + !! Uses Gaussian elimination, so will have issues if system is ill-conditioned. + !! Uses quad precision intermidiate values, so works best on small arrays. + use, intrinsic :: ieee_exceptions + implicit none + ! Arguments + integer(I4B), intent(in) :: n + real(DP), dimension(:,:), intent(in) :: A + real(DP), dimension(:), intent(in) :: b + logical, intent(out) :: lerr + ! Result + real(DP), dimension(n) :: x + ! Internals + real(QP), dimension(:), allocatable :: qx + type(ieee_status_type) :: original_fpe_status + logical, dimension(:), allocatable :: fpe_flag + + call ieee_get_status(original_fpe_status) ! Save the original floating point exception status + call ieee_set_flag(ieee_all, .false.) ! Set all flags to quiet + allocate(fpe_flag(size(ieee_usual))) + + qx = solve_wbs(ge_wpp(real(A, kind=QP), real(b, kind=QP))) + + call ieee_get_flag(ieee_usual, fpe_flag) + lerr = any(fpe_flag) + if (lerr .or. (any(abs(qx) > huge(x))) .or. (any(abs(qx) < tiny(x)))) then + x = 0.0_DP + else + x = real(qx, kind=DP) + end if + call ieee_set_status(original_fpe_status) + + return + end function util_solve_linear_system_d + + module function util_solve_linear_system_q(A,b,n,lerr) result(x) + !! Author: David A. Minton + !! + !! Solves the linear equation of the form A*x = b for x. + !! A is an (n,n) arrays + !! x and b are (n) arrays + !! Uses Gaussian elimination, so will have issues if system is ill-conditioned. + !! Uses quad precision intermidiate values, so works best on small arrays. + use, intrinsic :: ieee_exceptions + use swiftest + implicit none + ! Arguments + integer(I4B), intent(in) :: n + real(QP), dimension(:,:), intent(in) :: A + real(QP), dimension(:), intent(in) :: b + logical, intent(out) :: lerr + ! Result + real(QP), dimension(n) :: x + ! Internals + type(ieee_status_type) :: original_fpe_status + logical, dimension(:), allocatable :: fpe_flag + + call ieee_get_status(original_fpe_status) ! Save the original floating point exception status + call ieee_set_flag(ieee_all, .false.) ! Set all flags to quiet + allocate(fpe_flag(size(ieee_usual))) + + x = solve_wbs(ge_wpp(A, b)) + + call ieee_get_flag(ieee_usual, fpe_flag) + lerr = any(fpe_flag) + if (lerr) x = 0.0_DP + call ieee_set_status(original_fpe_status) + + return + end function util_solve_linear_system_q + + function solve_wbs(u) result(x) ! solve with backward substitution + !! Based on code available on Rosetta Code: https://rosettacode.org/wiki/Gaussian_elimination#Fortran + use, intrinsic :: ieee_exceptions + use swiftest + implicit none + ! Arguments + real(QP), intent(in), dimension(:,:), allocatable :: u + ! Result + real(QP), dimension(:), allocatable :: x + ! Internals + integer(I4B) :: i,n + + n = size(u, 1) + if (allocated(x)) deallocate(x) + if (.not.allocated(x)) allocate(x(n)) + if (any(abs(u) < tiny(1._DP)) .or. any(abs(u) > huge(1._DP))) then + x(:) = 0._DP + return + end if + call ieee_set_halting_mode(ieee_divide_by_zero, .false.) + do i = n, 1, -1 + x(i) = (u(i, n + 1) - sum(u(i, i + 1:n) * x(i + 1:n))) / u(i, i) + end do + return + end function solve_wbs + + function ge_wpp(A, b) result(u) ! gaussian eliminate with partial pivoting + !! Solve Ax=b using Gaussian elimination then backwards substitution. + !! A being an n by n matrix. + !! x and b are n by 1 vectors. + !! Based on code available on Rosetta Code: https://rosettacode.org/wiki/Gaussian_elimination#Fortran + use, intrinsic :: ieee_exceptions + use swiftest + implicit none + ! Arguments + real(QP), dimension(:,:), intent(in) :: A + real(QP), dimension(:), intent(in) :: b + ! Result + real(QP), dimension(:,:), allocatable :: u + ! Internals + integer(I4B) :: i,j,n,p + real(QP) :: upi + + n = size(a, 1) + allocate(u(n, (n + 1))) + u = reshape([A, b], [n, n + 1]) + call ieee_set_halting_mode(ieee_divide_by_zero, .false.) + do j = 1, n + p = maxloc(abs(u(j:n, j)), 1) + j - 1 ! maxloc returns indices between (1, n - j + 1) + if (p /= j) u([p, j], j) = u([j, p], j) + u(j + 1:, j) = u(j + 1:, j) / u(j, j) + do i = j + 1, n + 1 + upi = u(p, i) + if (p /= j) u([p, j], i) = u([j, p], i) + u(j + 1:n, i) = u(j + 1:n, i) - upi * u(j + 1:n, j) + end do + end do + return + end function ge_wpp + + module function util_solve_rkf45(f, y0in, t1, dt0, tol) result(y1) !! author: David A. Minton !! !! Implements the 4th order Runge-Kutta-Fehlberg ODE solver for initial value problems of the form f=dy/dt, y0 = y(t=0), solving for y1 = y(t=t1). Uses a 5th order adaptive step size control.