diff --git a/openmc/_utils.py b/openmc/_utils.py index b3c6d065c7..d00e4c58b7 100644 --- a/openmc/_utils.py +++ b/openmc/_utils.py @@ -31,31 +31,32 @@ def download(url, checksum=None, as_browser=False, **kwargs): page = Request(url, headers={'User-Agent': 'Mozilla/5.0'}) else: page = url - req = urlopen(page, **kwargs) - # Get file size from header - file_size = req.length - # Check if file already downloaded - basename = Path(urlparse(url).path).name - if os.path.exists(basename): - if os.path.getsize(basename) == file_size: - print('Skipping {}, already downloaded'.format(basename)) - return basename + with urlopen(page, **kwargs) as response: + # Get file size from header + file_size = response.length - # Copy file to disk in chunks - print('Downloading {}... '.format(basename), end='') - downloaded = 0 - with open(basename, 'wb') as fh: - while True: - chunk = req.read(_BLOCK_SIZE) - if not chunk: - break - fh.write(chunk) - downloaded += len(chunk) - status = '{:10} [{:3.2f}%]'.format( - downloaded, downloaded * 100. / file_size) - print(status + '\b'*len(status), end='') - print('') + # Check if file already downloaded + basename = Path(urlparse(url).path).name + if os.path.exists(basename): + if os.path.getsize(basename) == file_size: + print('Skipping {}, already downloaded'.format(basename)) + return basename + + # Copy file to disk in chunks + print('Downloading {}... '.format(basename), end='') + downloaded = 0 + with open(basename, 'wb') as fh: + while True: + chunk = response.read(_BLOCK_SIZE) + if not chunk: + break + fh.write(chunk) + downloaded += len(chunk) + status = '{:10} [{:3.2f}%]'.format( + downloaded, downloaded * 100. / file_size) + print(status + '\b'*len(status), end='', flush=True) + print('') if checksum is not None: downloadsum = hashlib.md5(open(basename, 'rb').read()).hexdigest() diff --git a/scripts/openmc-get-photon-data b/scripts/openmc-get-photon-data index 84660896e5..2d010b02ec 100755 --- a/scripts/openmc-get-photon-data +++ b/scripts/openmc-get-photon-data @@ -6,21 +6,20 @@ relaxation data and convert it to an HDF5 library for use with OpenMC. This data is used for photon transport in OpenMC. """ -import os -import sys -import shutil -import zipfile import argparse -from io import BytesIO -from urllib.request import urlopen +import os +from pathlib import Path +import zipfile import openmc.data +from openmc._utils import download class CustomFormatter(argparse.ArgumentDefaultsHelpFormatter, argparse.RawDescriptionHelpFormatter): pass + parser = argparse.ArgumentParser( description=__doc__, formatter_class=CustomFormatter @@ -36,47 +35,17 @@ block_size = 16384 # ============================================================================== # DOWNLOAD FILES FROM NNDC SITE -if not os.path.exists('photon_hdf5'): - os.mkdir('photon_hdf5') +output = Path('photon_hdf5') +output.mkdir(exist_ok=True) for f in files: - # Establish connection to URL - url = base_url + f - req = urlopen(url) - - # Get file size from header - file_size = req.length - downloaded = 0 - - # Check if file already downloaded - if os.path.exists(f): - if os.path.getsize(f) == file_size: - print('Skipping ' + f) - continue - else: - overwrite = input('Overwrite {}? ([y]/n) '.format(f)) - if overwrite.lower().startswith('n'): - continue - - # Copy file to disk - print('Downloading {}... '.format(f), end='') - with open(f, 'wb') as fh: - while True: - chunk = req.read(block_size) - if not chunk: - break - fh.write(chunk) - downloaded += len(chunk) - status = '{0:10} [{1:3.2f}%]'.format( - downloaded, downloaded * 100. / file_size) - print(status + chr(8)*len(status), end='') - print('') + download(base_url + f) # ============================================================================== # EXTRACT FILES for f in files: - print('Extracting {0}...'.format(f)) + print('Extracting {}...'.format(f)) zipfile.ZipFile(f).extractall() # ============================================================================== @@ -87,9 +56,10 @@ if args.cross_sections is not None: lib_path = args.cross_sections library = openmc.data.DataLibrary.from_xml(lib_path) else: - lib_path = os.path.join('photon_hdf5', 'cross_sections.xml') + lib_path = output / 'cross_sections.xml' library = openmc.data.DataLibrary() +# Iterate over each natural element from Z=1 to Z=100 for z in range(1, 101): element = openmc.data.ATOMIC_SYMBOL[z] print('Generating HDF5 file for Z={} ({})...'.format(z, element)) @@ -97,11 +67,11 @@ for z in range(1, 101): # Generate instance of IncidentPhoton photo_file = os.path.join('photoat', 'photoat-{:03}_{}_000.endf'.format(z, element)) atom_file = os.path.join('atomic_relax', 'atom-{:03}_{}_000.endf'.format(z, element)) - f = openmc.data.IncidentPhoton.from_endf(photo_file, atom_file) + data = openmc.data.IncidentPhoton.from_endf(photo_file, atom_file) # Write HDF5 file and register it - hdf5_file = os.path.join('photon_hdf5', element + '.h5') - f.export_to_hdf5(hdf5_file, 'w') + hdf5_file = output / (element + '.h5') + data.export_to_hdf5(hdf5_file, 'w') library.register_file(hdf5_file) library.export_to_xml(lib_path)