Merge pull request pyne#10 from Shimwell/jeff32_temperature_Control

Jeff32 temperature control

convert_jeff32.py

@@ -6,6 +6,7 @@
 import tarfile
 import zipfile
 from collections import defaultdict
+from pathlib import Path
 from string import digits
 from urllib.parse import urljoin
 
@@ -18,12 +19,6 @@
 
 """
 
-download_warning = """
-WARNING: This script will download approximately 9 GB of data. Extracting and
-processing the data may require as much as 40 GB of additional free disk
-space. Note that if you don't need all 11 temperatures, you can modify the
-'files' list in the script to download only the data you want.
-"""
 
 class CustomFormatter(argparse.ArgumentDefaultsHelpFormatter,
                       argparse.RawDescriptionHelpFormatter):
@@ -33,7 +28,7 @@ class CustomFormatter(argparse.ArgumentDefaultsHelpFormatter,
     description=description,
     formatter_class=CustomFormatter
 )
-parser.add_argument('-d', '--destination', default='jeff32_hdf5',
+parser.add_argument('-d', '--destination', type=Path, default=None,
                     help='Directory to create new library in')
 parser.add_argument('--download', action='store_true',
                     help='Download files from OECD-NEA')
@@ -47,59 +42,83 @@ class CustomFormatter(argparse.ArgumentDefaultsHelpFormatter,
                     default='latest', help="Output HDF5 versioning. Use "
                     "'earliest' for backwards compatibility or 'latest' for "
                     "performance")
+parser.add_argument('-r', '--release', choices=['3.2'],
+                    default='3.2', help="The nuclear data library release version. "
+                    "The currently supported options are 3.2")
+parser.add_argument('-t', '--temperatures', 
+                    choices=['293', '400', '500', '600', '700', '800', '900', 
+                             '1000', '1200', '1500', '1800'],
+                    default=['293', '400', '500', '600', '700', '800', '900',
+                             '1000', '1200', '1500', '1800'], 
+                    help="Temperatures to download in Kelvin", nargs='+',)    
 parser.set_defaults(download=True, extract=True)
 args = parser.parse_args()
 
-print(download_warning)
-
-base_url = 'https://www.oecd-nea.org/dbforms/data/eva/evatapes/jeff_32/Processed/'
-files = ['JEFF32-ACE-293K.tar.gz',
-         'JEFF32-ACE-400K.tar.gz',
-         'JEFF32-ACE-500K.tar.gz',
-         'JEFF32-ACE-600K.tar.gz',
-         'JEFF32-ACE-700K.tar.gz',
-         'JEFF32-ACE-800K.zip',
-         'JEFF32-ACE-900K.tar.gz',
-         'JEFF32-ACE-1000K.tar.gz',
-         'JEFF32-ACE-1200K.tar.gz',
-         'JEFF32-ACE-1500K.tar.gz',
-         'JEFF32-ACE-1800K.tar.gz',
-         'TSLs.tar.gz']
+library_name = 'jeff'
+ace_files_dir = Path('-'.join([library_name, args.release, 'ace']))
+# the destination is decided after the release is know to avoid putting the release in a folder with a misleading name
+if args.destination is None:
+    args.destination = Path('-'.join([library_name, args.release, 'hdf5']))
+
+# This dictionary contains all the unique information about each release. This can be exstened to accommodated new releases
+release_details = {
+    '3.2':{
+        'base_url': 'https://www.oecd-nea.org/dbforms/data/eva/evatapes/jeff_32/Processed/',
+        'files': [f'JEFF32-ACE-{t}K.zip' if t=='800' else f'JEFF32-ACE-{t}K.tar.gz' for t in args.temperatures] 
+                  +['TSLs.tar.gz'],
+        'neutron_files':ace_files_dir.rglob('*.ACE'),
+        'metastables': ace_files_dir.rglob('*M.ACE'),
+        'sab_files': ace_files_dir.glob('ANNEX_6_3_STLs/*/*.ace'),
+        'redundant': ace_files_dir.glob('ACEs_293K/*-293.ACE'),
+        'compressed_file_size': '9 GB',
+        'uncompressed_file_size': '40 GB'
+    }
+}
+
+download_warning = """
+WARNING: This script will download up to {} GB of data. Extracting and
+processing the data may require as much as {} GB of additional free disk
+space. Note that if you don't need all 11 temperatures, you can used the 
+--temperature argument to download only the temperatures you want.
+""".format(release_details[args.release]['compressed_file_size'],
+           release_details[args.release]['uncompressed_file_size'])
 
 # ==============================================================================
 # DOWNLOAD FILES FROM OECD SITE
 
 if args.download:
-    for f in files:
-        download(urljoin(base_url, f))
+    print(download_warning)
+    for f in release_details[args.release]['files']:
+        download(urljoin(release_details[args.release]['base_url'], f))
 
 # ==============================================================================
 # EXTRACT FILES FROM TGZ
 
 if args.extract:
-    for f in files:
+    for f in release_details[args.release]['files']:
         # Extract files
         if f.endswith('.zip'):
             with zipfile.ZipFile(f, 'r') as zipf:
                 print('Extracting {}...'.format(f))
-                zipf.extractall('jeff-3.2')
+                zipf.extractall(ace_files_dir)
 
         else:
             suffix = 'ACEs_293K' if '293' in f else ''
             with tarfile.open(f, 'r') as tgz:
                 print('Extracting {}...'.format(f))
-                tgz.extractall(os.path.join('jeff-3.2', suffix))
+                tgz.extractall(ace_files_dir / suffix)
 
             # Remove thermal scattering tables from 293K data since they are
             # redundant
             if '293' in f:
-                for path in glob.glob(os.path.join('jeff-3.2', 'ACEs_293K', '*-293.ACE')):
+                for path in release_details[args.release]['redundant']:
+                    print(f'removing {path}')
                     os.remove(path)
 
 # ==============================================================================
 # CHANGE ZAID FOR METASTABLES
 
-metastables = glob.glob(os.path.join('jeff-3.2', '**', '*M.ACE'))
+metastables = release_details[args.release]['metastables']
 for path in metastables:
     print('    Fixing {} (ensure metastable)...'.format(path))
     text = open(path, 'r').read()
@@ -112,23 +131,21 @@ class CustomFormatter(argparse.ArgumentDefaultsHelpFormatter,
 # GENERATE HDF5 LIBRARY -- NEUTRON FILES
 
 # Get a list of all ACE files
-neutron_files = glob.glob(os.path.join('jeff-3.2', '*', '*.ACE'))
+neutron_files = release_details[args.release]['neutron_files']
 
 # Group together tables for same nuclide
 tables = defaultdict(list)
 for filename in sorted(neutron_files):
-    dirname, basename = os.path.split(filename)
-    name = basename.split('.')[0]
-    tables[name].append(filename)
+    name = filename.stem
+    tables[name].append(str(filename))
 
 # Sort temperatures from lowest to highest
 for name, filenames in sorted(tables.items()):
     filenames.sort(key=lambda x: int(
         x.split(os.path.sep)[1].split('_')[1][:-1]))
 
 # Create output directory if it doesn't exist
-if not os.path.isdir(args.destination):
-    os.mkdir(args.destination)
+args.destination.mkdir(parents=True, exist_ok=True)
 
 library = openmc.data.DataLibrary()
 
@@ -143,7 +160,7 @@ class CustomFormatter(argparse.ArgumentDefaultsHelpFormatter,
         data.add_temperature_from_ace(filename)
 
     # Export HDF5 file
-    h5_file = os.path.join(args.destination, data.name + '.h5')
+    h5_file = args.destination / f'{data.name}.h5'
     print('Writing {}...'.format(h5_file))
     data.export_to_hdf5(h5_file, 'w', libver=args.libver)
 
@@ -153,14 +170,12 @@ class CustomFormatter(argparse.ArgumentDefaultsHelpFormatter,
 # ==============================================================================
 # GENERATE HDF5 LIBRARY -- S(A,B) FILES
 
-sab_files = glob.glob(os.path.join('jeff-3.2', 'ANNEX_6_3_STLs', '*', '*.ace'))
-
 # Group together tables for same nuclide
 tables = defaultdict(list)
-for filename in sorted(sab_files):
+for filename in sorted(release_details[args.release]['sab_files']):
     dirname, basename = os.path.split(filename)
     name = basename.split('-')[0]
-    tables[name].append(filename)
+    tables[name].append(str(filename))
 
 # Sort temperatures from lowest to highest
 for name, filenames in sorted(tables.items()):
@@ -186,13 +201,12 @@ class CustomFormatter(argparse.ArgumentDefaultsHelpFormatter,
         data.add_temperature_from_ace(table)
 
     # Export HDF5 file
-    h5_file = os.path.join(args.destination, data.name + '.h5')
+    h5_file = args.destination / f'{data.name}.h5'
     print('Writing {}...'.format(h5_file))
     data.export_to_hdf5(h5_file, 'w', libver=args.libver)
 
     # Register with library
     library.register_file(h5_file)
 
 # Write cross_sections.xml
-libpath = os.path.join(args.destination, 'cross_sections.xml')
-library.export_to_xml(libpath)
+library.export_to_xml(args.destination / 'cross_sections.xml')