added a radbelt_type with a get_flux method

test updates and added a python test

src/core.f90

@@ -14,19 +14,27 @@ module core
 
    implicit none
 
-   public :: get_flux
+   type,public :: radbelt_type
+      !! the main class that can be used to get the flux.
+      private
+      type(trm_type) :: trm
+      type(shellig_type) :: igrf
+      contains
+      private
+      procedure,public :: get_flux => get_flux_
+   end type radbelt_type
+
+   public :: get_flux !! simple function version for testing
 
    contains
 
 !*****************************************************************************************
 !>
 !  Calculate the flux of trapped particles at a specific location and time.
-!
-!@note This routine is not efficient since it will reload all the
-!      files every time it is called.
 
-   function get_flux(Lon,Lat,Height,Year,E,Imname) result(flux)
+   function get_flux_(me,lon,lat,height,year,e,imname) result(flux)
 
+      class(radbelt_type),intent(inout) :: me
       real(wp),intent(in) :: lon !! geodetic longitude in degrees (east)
       real(wp),intent(in) :: lat !! geodetic latitude in degrees (north)
       real(wp),intent(in) :: height !! altitude in km above sea level
@@ -43,11 +51,40 @@ function get_flux(Lon,Lat,Height,Year,E,Imname) result(flux)
 
       real(wp) :: xl !! l value
       real(wp) :: bbx
-      type(trm_type) :: trm
-      type(shellig_type) :: igrf
 
-      call igrf%igrf(lon,lat,height,year,xl,bbx)
-      call trm%aep8(e,xl,bbx,imname,flux)
+      call me%igrf%igrf(lon,lat,height,year,xl,bbx)
+      call me%trm%aep8(e,xl,bbx,imname,flux)
+
+   end function get_flux_
+!*****************************************************************************************
+
+!*****************************************************************************************
+!>
+!  Calculate the flux of trapped particles at a specific location and time.
+!  This is just a function version of the class method from [[radbelt_type]].
+!
+!@note This routine is not efficient at all since it will reload all the
+!      files every time it is called.
+
+   function get_flux(lon,lat,height,year,e,imname) result(flux)
+
+      real(wp),intent(in) :: lon !! geodetic longitude in degrees (east)
+      real(wp),intent(in) :: lat !! geodetic latitude in degrees (north)
+      real(wp),intent(in) :: height !! altitude in km above sea level
+      real(wp),intent(in) :: year !! decimal year for which geomagnetic field is to
+                                  !! be calculated (e.g.:1995.5 for day 185 of 1995)
+      real(wp),intent(in) :: e !! minimum energy
+      integer,intent(in) :: imname !! which method to use:
+                                   !!
+                                   !! * 1 -- particle species: electrons, solar activity: min
+                                   !! * 2 -- particle species: electrons, solar activity: max
+                                   !! * 3 -- particle species: protons, solar activity: min
+                                   !! * 4 -- particle species: protons, solar activity: max
+      real(wp) :: flux !! The flux of particles above the given energy, in units of cm^-2 s^-1.
+
+      type(radbelt_type) :: radbelt
+
+      flux = radbelt%get_flux(lon,lat,height,year,e,imname)
 
    end function get_flux
 !*****************************************************************************************

src/shellig.f90

@@ -53,6 +53,9 @@ module shellig_module
          integer :: nmax = 0 !! maximum order of spherical harmonics
          real(wp) :: Time = 0.0_wp !! year (decimal: 1973.5) for which magnetic field is to be calculated
          real(wp),dimension(144) :: g = 0.0_wp  !! `g(m)` -- normalized field coefficients (see [[feldcof]]) m=nmax*(nmax+2)
+         integer :: nmax1 = 0 !! saved variables from the file
+         integer :: nmax2 = 0 !! saved variables from the file
+         real(wp),dimension(144) :: g_cache = 0.0_wp  !! saved `g` from the file
 
          ! formerly saved vars in shellg:
          real(wp) :: step = 0.20_wp !! step size for field line tracing
@@ -99,6 +102,13 @@ subroutine igrf(me,lon,lat,height,year,xl,bbx)
 
       real(wp),parameter :: stps = 0.05_wp
 
+      ! JW : do we need to reset some or all of these ?
+      me%sp = 0.0_wp
+      me%xi = 0.0_wp
+      me%h = 0.0_wp
+      me%step = 0.20_wp
+      me%steq = 0.03_wp
+
       call me%feldcof(year,dimo)
       call me%feldg(lat,lon,height,bnorth,beast,bdown,babs)
       call me%shellg(lat,lon,height,dimo,xl,icode,bab1)
@@ -651,7 +661,7 @@ subroutine feldcof(me,year,dimo)
                                 !! to earth's radius) at the time (year)
 
    real(wp) :: dte1 , dte2 , erad , gha(144) , sqrt2
-   integer :: i , ier , j , l , m , n , nmax1 , nmax2, iyea
+   integer :: i , ier , j , l , m , n , iyea
    character(len=filename_len) :: fil2
    real(wp) :: x , f0 , f !! these were double precision in original
                           !! code while everything else was single precision
@@ -690,16 +700,19 @@ subroutine feldcof(me,year,dimo)
    if (read_file) then
       ! get igrf coefficients for the boundary years
       ! [if they have not ready been loaded]
-      call me%getshc(me%name,nmax1,erad,me%g,ier)
+      call me%getshc(me%name,me%nmax1,erad,me%g,ier)
       if ( ier/=0 ) error stop 'error reading file: '//trim(me%name)
-      call me%getshc(fil2,nmax2,erad,me%gh2,ier)
+      me%g_cache = me%g ! because it is modified below, we have to cache the original values from the file
+      call me%getshc(fil2,me%nmax2,erad,me%gh2,ier)
       if ( ier/=0 ) error stop 'error reading file: '//trim(fil2)
+   else
+      me%g = me%g_cache
    end if
    !-- determine igrf coefficients for year
    if ( l<=numye-1 ) then
-      call me%intershc(year,dte1,nmax1,me%g,dte2,nmax2,me%gh2,me%nmax,gha)
+      call me%intershc(year,dte1,me%nmax1,me%g,dte2,me%nmax2,me%gh2,me%nmax,gha)
    else
-      call me%extrashc(year,dte1,nmax1,me%g,nmax2,me%gh2,me%nmax,gha)
+      call me%extrashc(year,dte1,me%nmax1,me%g,me%nmax2,me%gh2,me%nmax,gha)
    endif
    !-- determine magnetic dipol moment and coeffiecients g
    f0 = 0.0_wp

src/trmfun.f90

@@ -64,6 +64,11 @@ subroutine aep8(me,e,l,bb0,imname,flux)
 
       name = mname(Imname) ! the file to load
 
+      ! JW : do we need to reset some or all of these ?
+      me%fistep = 0.0_wp
+      me%f1 = 1.001_wp
+      me%f2 = 1.002_wp
+
       ! check to see if this file has already been loaded
       ! [the class can store one file at a time]
       load_file = .true.

test/radbelt_test.f90

@@ -8,7 +8,10 @@ program radbelt_test
     implicit none
 
     real(wp) :: lon, lat, height, year, e, flux, error, relerror
-    integer :: imname, i
+    integer :: imname, i, ilat, ilon, ialt
+    real(wp) :: tstart, tend
+    type(radbelt_type) :: radbelt
+    integer :: n_cases
 
     lon = -45.0_wp
     lat = -30.0_wp
@@ -22,19 +25,46 @@ program radbelt_test
     Imname = 4 ! 'p', 'max'
     e = 20.0_wp
 
-    do i = 1, 3
+    ! simple example:
     flux = get_flux(lon,lat,height,year,e,imname)
-
     ! error = Flux - 2642.50268555_wp  ! difference from python wrapper version (radbelt)
     error = Flux - 2642.50370051985726336559603128948869_wp ! difference from real128 version
     relerror = abs(error/flux)
-
     write(*,*) 'Flux      = ', flux
     write(*,*) 'Error     = ', error
     write(*,*) 'Rel Error = ', relerror
-
     if (relerror>10*epsilon(1.0_wp)) error stop 'error'
+    write(*,*) ''
 
+    ! speed tests:
+    call cpu_time(tstart)
+    n_cases = 0
+    do ilat = -89, 90, 5
+        do ilon = -180, 180, 45
+            do ialt = 500, 1000, 100
+                n_cases = n_cases + 1
+                flux = get_flux(real(ilon,wp),real(ilat,wp),real(ialt,wp),year,e,imname)
+                !if (flux/=0.0_wp) write(*,*) ilat, ilon, ialt, flux
+            end do
+        end do
+    end do
+    call cpu_time(tend)
+    write(*,'(a25,f6.3,a,i7,a)') 'Function version runtime: ', (tend - tstart), ' sec. ', &
+                                 int(n_cases/(tend - tstart)), ' (cases/sec)'
+    call cpu_time(tstart)
+    n_cases = 0
+    do ilat = -90, 90, 5
+        do ilon = -180, 180, 45
+            do ialt = 500, 1000, 100
+                n_cases = n_cases + 1
+                flux = radbelt%get_flux(real(ilon,wp),real(ilat,wp),real(ialt,wp),year,e,imname)
+                !write(*,*) ilat, ilon, ialt, flux
+            end do
+        end do
     end do
+    call cpu_time(tend)
+    write(*,'(a25,f6.3,a,i7,a)') 'Class version runtime: ', (tend - tstart), ' sec. ', &
+                                 int(n_cases/(tend - tstart)), ' (cases/sec)'
 
+    write(*,*) 'n_cases = ', n_cases
 end program radbelt_test

test/test.py

@@ -0,0 +1,30 @@
+#
+# Test of the Python version from: https://github.com/nasa/radbelt
+#
+
+from radbelt import get_flux
+from astropy import units as u
+from astropy.coordinates import EarthLocation
+from astropy.time import Time
+import time
+
+t = Time('2021-03-01')
+energy = 20 * u.MeV
+
+tstart = time.time()
+n_cases = 0
+
+# fortran loop indices:
+# do ilat = -90, 90, 5
+#     do ilon = -180, 180, 45
+#         do ialt = 500, 1000, 100
+
+for lat in range(-90, 91, 5):
+    for lon in range(-180, 181, 45):
+        for alt in range(500, 1001, 100):
+            n_cases = n_cases + 1
+            coords = EarthLocation(lon * u.deg, lat * u.deg, alt * u.km)
+            f = get_flux(coords, t, energy, 'p', 'max')  # doctest: +FLOAT_CMP
+            #print(t.utc.decimalyear, lat, lon, alt, f.value)
+tend = time.time()
+print(f'Python version runtime: {tend-tstart} sec. {int(n_cases/(tend-tstart))} (cases/sec)')