Skip to content

Commit

Permalink
Use download() in openmc-get-photon-data and respond to other comments
Browse files Browse the repository at this point in the history
  • Loading branch information
paulromano committed Jan 27, 2020
1 parent 04308b7 commit 37af057
Show file tree
Hide file tree
Showing 2 changed files with 37 additions and 66 deletions.
47 changes: 24 additions & 23 deletions openmc/_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
56 changes: 13 additions & 43 deletions scripts/openmc-get-photon-data
Original file line number Diff line number Diff line change
Expand Up @@ -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 argparse
import os
import sys
import shutil
from pathlib import Path
import zipfile
import argparse
from io import BytesIO
from urllib.request import urlopen

import openmc.data
from openmc._utils import download


class CustomFormatter(argparse.ArgumentDefaultsHelpFormatter,
argparse.RawDescriptionHelpFormatter):
pass


parser = argparse.ArgumentParser(
description=__doc__,
formatter_class=CustomFormatter
Expand All @@ -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()

# ==============================================================================
Expand All @@ -87,21 +56,22 @@ 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))

# 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)

0 comments on commit 37af057

Please sign in to comment.