Source code for pydrad.configure.configure
"""
Configure HYDRAD simulations
"""
import asdf
import astropy.units as u
import copy
import datetime
import numpy as np
import os
import pathlib
import shutil
import stat
import tempfile
from jinja2 import ChoiceLoader, DictLoader, Environment, PackageLoader
from pydrad.configure import filters
from pydrad.configure.util import get_equilibrium_heating_rate, run_shell_command
__all__ = ['Configure']
[docs]
class Configure:
"""
Configure HYDRAD simulations from a single Python `dict`
Parameters
----------
config : `dict`
All input parameters for configuring simulation
"""
def __init__(self, config, **kwargs):
self.config = copy.deepcopy(config)
loader = ChoiceLoader([
DictLoader(kwargs.get('templates', {})),
PackageLoader('pydrad', 'configure/templates')
])
self.env = Environment(loader=loader)
self.env.filters['units_filter'] = filters.units_filter
self.env.filters['log10_filter'] = filters.log10_filter
self.env.filters['get_atomic_symbol'] = filters.get_atomic_symbol
self.env.filters['get_atomic_number'] = filters.get_atomic_number
self.env.filters['sort_elements'] = filters.sort_elements
self.env.filters['is_required'] = filters.is_required
self.env.filters['sci_notation'] = filters.sci_notation
# Compiler flags
self.optimization_flags = kwargs.get('optimization_flags')
self.compiler = kwargs.get('compiler', 'g++')
# NOTE: Freeze the date at instantiation so that files can be compared
# exactly for testing
if kwargs.get('freeze_date', False):
self._date = self.date
self._freeze_date = True
[docs]
@staticmethod
def load_config(filename):
"""
Load a base configuration from an ASDF file
Parameters
----------
filename : `str`
"""
with asdf.open(filename) as af:
config = copy.deepcopy(dict(af.tree))
return config
[docs]
def save_config(self, filename):
"""
Save the simulation configuration as an ASDF file
Parameters
----------
filename : `str`
"""
asdf.AsdfFile(self.config).write_to(filename)
[docs]
def setup_simulation(self,
output_path,
base_path,
run_initial_conditions=True,
overwrite=False,
**kwargs):
"""
Setup a HYDRAD simulation with desired outputs from a clean copy
Parameters
----------
output_path : `str`
Path to new copy of HYDRAD
base_path : `str`
Path to existing HYDRAD
run_initial_conditions : `bool`
If True, compile and run the initial conditions code
overwrite : `bool`
If True, overwrite an existing HYDRAD instance at ``output_path``
if it exists.
"""
with tempfile.TemporaryDirectory() as tmpdir:
# NOTE: this is all done in a temp directory and then copied over
# so that if something fails, all the files are cleaned up
shutil.copytree(base_path, tmpdir, dirs_exist_ok=True)
if run_initial_conditions:
execute = kwargs.get('execute', True)
self.setup_initial_conditions(tmpdir, execute=execute)
self.setup_hydrad(tmpdir)
self.save_config(os.path.join(tmpdir, 'pydrad_config.asdf'))
shutil.copytree(tmpdir, output_path, dirs_exist_ok=overwrite)
[docs]
def setup_initial_conditions(self, root_dir, execute=True):
"""
Compile and run the initial conditions code to get the initial
loop profile
Parameters
----------
root_dir : path-like
Path to new HYDRAD copy
execute : `bool`
If True (default), compute initial conditions. Otherwise, they
are only compiled. This is useful for debugging.
"""
root_dir = pathlib.Path(root_dir)
build_script_filename = pathlib.Path('Initial_Conditions') / 'build_scripts' / 'build_script.bat'
files = [
('Initial_Conditions/source/config.h', self.initial_conditions_header),
('Initial_Conditions/config/initial_conditions.cfg', self.initial_conditions_cfg),
('Radiation_Model/source/config.h', self.radiation_header),
('Radiation_Model/config/elements_eq.cfg', self.radiation_equilibrium_cfg),
('Radiation_Model/config/elements_neq.cfg', self.radiation_nonequilibrium_cfg),
(build_script_filename, self.initial_conditions_build_script),
]
# NOTE: there are two options here so that the gravitational and
# magnetic field polynomial fits can be applied just to the
# hydrodynamic step and not to the initial conditions. Sometimes an
# initial condition cannot be found if the gravitational and/or magnetic
# field profile is a bit strange.
if (self.config['initial_conditions']['use_poly_fit_gravity']
and 'poly_fit_gravity' in self.config['general']):
files += [('poly_fit.gravity', self.poly_fit_gravity)]
if (self.config['initial_conditions']['use_poly_fit_magnetic_field']
and 'poly_fit_magnetic_field' in self.config['general']):
files += [('poly_fit.magnetic_field', self.poly_fit_magnetic_field)]
for filename, filestring in files:
with (root_dir / filename).open(mode='w') as f:
f.write(filestring)
# NOTE: make sure we have needed permissions to run compile script
os.chmod(root_dir / build_script_filename, mode=stat.S_IRWXU)
run_shell_command(root_dir / build_script_filename)
(root_dir / 'Initial_Conditions' / 'profiles').mkdir(parents=True, exist_ok=True)
if execute:
run_shell_command(root_dir / 'Initial_Conditions.exe')
if self.config['heating']['background'].get('use_initial_conditions', False):
self.equilibrium_heating_rate = get_equilibrium_heating_rate(root_dir)
[docs]
def setup_hydrad(self, root_dir):
"""
Compile HYDRAD code with appropriate header and config files
Parameters
-----------
root_dir : `str`
Path to new HYDRAD copy
"""
root_dir = pathlib.Path(root_dir)
build_script_filename = pathlib.Path('HYDRAD') / 'build_scripts' / 'build_script.bat'
files = [
('Radiation_Model/source/config.h', self.radiation_header),
('Radiation_Model/config/elements_eq.cfg', self.radiation_equilibrium_cfg),
('Radiation_Model/config/elements_neq.cfg', self.radiation_nonequilibrium_cfg),
('Heating_Model/source/config.h', self.heating_header),
('Heating_Model/config/heating_model.cfg', self.heating_cfg),
('HYDRAD/source/config.h', self.hydrad_header),
('HYDRAD/source/collisions.h', self.collisions_header),
('HYDRAD/config/hydrad.cfg', self.hydrad_cfg),
(build_script_filename, self.hydrad_build_script),
]
if 'poly_fit_gravity' in self.config['general']:
files += [('poly_fit.gravity', self.poly_fit_gravity)]
if 'poly_fit_magnetic_field' in self.config['general']:
files += [('poly_fit.magnetic_field', self.poly_fit_magnetic_field)]
if self.config['heating'].get('beam', False):
files += [('Heating_Model/config/beam_heating_model.cfg',
self.beam_heating_cfg)]
for filename, filestring in files:
with (root_dir / filename).open(mode='w') as f:
f.write(filestring)
# NOTE: make sure we have needed permissions to run compile script
os.chmod(root_dir / build_script_filename, mode=stat.S_IRWXU)
run_shell_command(root_dir / build_script_filename)
(root_dir / 'Results').mkdir(parents=True, exist_ok=True)
@property
def date(self):
"""
Return the current date
"""
# NOTE: freeze date at instantiation for testing purposes
if hasattr(self, '_freeze_date') and self._freeze_date:
return self._date
else:
return datetime.datetime.now(datetime.UTC).strftime('%Y-%m-%d_%H.%M.%S UTC')
@property
def templates(self,):
"""
List of available templates
"""
return self.env.list_templates()
[docs]
def get_raw_template(self, name):
"""
Return the unrendered template.
"""
with pathlib.Path(self.env.get_template(name).filename).open() as f:
return f.read()
@property
def initial_conditions_cfg(self):
"""
Initial conditions configuration file,
`Initial_Conditions/config/initial_conditions.cfg`
"""
return self.env.get_template('initial_conditions.cfg').render(
date=self.date,
**self.config
)
@property
def initial_conditions_header(self):
"""
Initial conditions header file, `Initial_Conditions/source/config.h`
"""
return self.env.get_template('initial_conditions.config.h').render(
date=self.date,
maximum_cells=self.maximum_cells,
minimum_cells=self.minimum_cells,
**self.config
)
@property
def hydrad_cfg(self):
"""
HYDRAD configuration file, `HYDRAD/config/hydrad.cfg`
"""
return self.env.get_template('hydrad.cfg').render(
date=self.date,
**self.config
)
@property
def hydrad_header(self):
"""
HYDRAD header file, `HYDRAD/source/config.h`
"""
return self.env.get_template('hydrad.config.h').render(
date=self.date,
**self.config
)
@property
def heating_cfg(self):
"""
Heating model configuration file, `Heating_Model/config/heating.cfg`.
If background heating is enabled and you want to use the equilibrium
values, you must run the initial conditions and set the
`equilibrium_heating_rate` attribute first.
"""
background = copy.deepcopy(self.config['heating']['background'])
if background.get('use_initial_conditions', False):
background['rate'] = self.equilibrium_heating_rate
background['location'] = self.config['initial_conditions']['heating_location']
background['scale_height'] = self.config['initial_conditions']['heating_scale_height']
return self.env.get_template('heating.cfg').render(
date=self.date,
background=background,
**self.config
)
@property
def heating_header(self):
"""
Heating model header file, `Heating_Model/source/config.h`
"""
return self.env.get_template('heating.config.h').render(
date=self.date,
**self.config
)
@property
def beam_heating_cfg(self):
"""
Beam heating model configuration file.
"""
return self.env.get_template('heating.beam.cfg').render(
date=self.date,
**self.config,
)
@property
def radiation_equilibrium_cfg(self):
"""
Equilibrium elements configuration file,
`Radiation_Model/config/elements_eq.cfg`
"""
elements = self.config['radiation'].get('elements_equilibrium', [])
return self.env.get_template('radiation.elements.cfg').render(
date=self.date,
elements=elements,
**self.config
)
@property
def radiation_nonequilibrium_cfg(self):
"""
Non-equilibrium elements configuration file,
`Radiation_Model/config/elements_neq.cfg`
"""
elements = self.config['radiation'].get('elements_nonequilibrium', [])
return self.env.get_template('radiation.elements.cfg').render(
date=self.date,
elements=elements,
**self.config
)
@property
def radiation_header(self):
"""
Radiation model header file, `Radiation_Model/source/config.h`
"""
return self.env.get_template('radiation.config.h').render(
date=self.date,
**self.config
)
@property
def collisions_header(self):
"""
Collisions header file, `HYDRAD/source/collisions.h`
"""
return self.env.get_template('collisions.h').render(
date=self.date,
**self.config
)
@property
def poly_fit_magnetic_field(self):
"""
Sixth-order polynomial fit coefficients for computing flux tube
expansion
"""
fit_results = self._fit_poly_domains(
**self.config['general']['poly_fit_magnetic_field'],
target_unit='G',
)
return self.env.get_template('coefficients.cfg').render(
date=self.date,
**fit_results,
)
@property
def poly_fit_gravity(self):
"""
Sixth-order polynomial fit coefficients for computing gravitational
acceleration
"""
fit_results = self._fit_poly_domains(
**self.config['general']['poly_fit_gravity'],
target_unit='cm s-2'
)
return self.env.get_template('coefficients.cfg').render(
date=self.date,
**fit_results,
)
def _fit_poly_domains(self, x, y, domains, order, target_unit):
"""
Perform polynomial fit to quantity as a function of field aligned coordinate
over multiple domains and return fitting coefficients.
"""
x = (x/self.config['general']['loop_length']).to_value(u.dimensionless_unscaled)
y = y.to_value(target_unit)
coefficients = []
minmax = []
for i in range(len(domains)-1):
i_d = np.where(np.logical_and(x>=domains[i], x<=domains[i+1]))
# NOTE: The order is reversed because HYDRAD expects the opposite order of
# what NumPy outputs
coefficients.append(np.polyfit(x[i_d], y[i_d], order)[::-1])
minmax.append([y[i_d].min(), y[i_d].max()])
return {
'domains': domains,
'order': order,
'minmax': minmax,
'coefficients': coefficients,
}
@property
def minimum_cells(self):
r"""
Minimum allowed number of grid cells,
:math:`n_{min}=\lceil L/\Delta s_{max}\rceil`, where :math:`L` is the loop
length and :math:`\Delta s_{max}` is the maximum allowed grid cell width.
Optionally, if the minimum number of cells is specified
in ``config['grid']['minimum_cells']``, this value will take
precedence.
"""
if 'minimum_cells' in self.config['grid']:
return int(self.config['grid']['minimum_cells'])
L = self.config['general']['loop_length']
n_min = L / self.config['grid']['maximum_cell_width']
if n_min.decompose().unit != u.dimensionless_unscaled:
raise u.UnitConversionError(f'''Maximum cell width must be able to be converted to {L.unit}''')
return int(np.ceil(n_min.decompose()))
@property
def maximum_cells(self):
r"""
Maximum allowed number of grid cells,
:math:`n_{max}=\lfloor 2^{L_R}n_{min}\rfloor`, where :math:`L_R` is the maximum
refinement level and :math:`n_{min}` is the minimum allowed number of
grid cells. Optionally, if the maximum number of cells is specified
in ``config['grid']['maximum_cells']``, this value will take
precedence. Alternatively, using a value of 30000 will be sufficiently large for
any amount of refinement used in HYDRAD.
"""
if 'maximum_cells' in self.config['grid']:
return int(self.config['grid']['maximum_cells'])
n_min = self.config['general']['loop_length'] / self.config['grid']['maximum_cell_width']
if n_min.decompose().unit != u.dimensionless_unscaled:
raise u.UnitConversionError(
f'''Maximum cell width must be able to be converted to
{self.config['general']['loop_length'].unit}''')
return int(np.floor(
2**self.config['grid']['maximum_refinement_level'] * n_min))
@property
def optimization_flags(self):
return self._optimization_flags
@optimization_flags.setter
def optimization_flags(self, value):
if value is None:
value = ['O3', 'flto', 'Wno-unused-variable', 'Wno-write-strings']
self._optimization_flags = [f'-{f}' for f in value]
@property
def initial_conditions_build_script(self):
files = [
'../source/main.cpp',
'../source/ode.cpp',
'../source/misc.cpp',
'../../Radiation_Model/source/element.cpp',
'../../Radiation_Model/source/radiation.cpp',
'../../Radiation_Model/source/OpticallyThick/OpticallyThickIon.cpp ',
'../../Radiation_Model/source/OpticallyThick/RadiativeRates.cpp ',
'../../Resources/source/gammabeta.cpp ',
'../../Resources/source/fitpoly.cpp ',
'../../Resources/Utils/generatePieceWiseFit/source/piecewisefit.cpp ',
'../../Resources/Utils/regPoly/regpoly.cpp ',
'../../Resources/Utils/regPoly/nrutil.cpp ',
'../../Resources/source/file.cpp',
]
return self.env.get_template('build_script.bat').render(
compiler=self.compiler,
files=files,
flags=['-Wall',] + self.optimization_flags,
executable='../../Initial_Conditions.exe',
)
@property
def hydrad_build_script(self):
files = [
'../source/main.cpp',
'../source/cell.cpp',
'../source/mesh.cpp',
'../source/eqns.cpp',
'../../Kinetic_Model/source/kinetic.cpp',
'../../Kinetic_Model/source/gamma.cpp',
'../../Heating_Model/source/heat.cpp',
'../../Radiation_Model/source/ionfrac.cpp',
'../../Radiation_Model/source/element.cpp',
'../../Radiation_Model/source/radiation.cpp',
'../../Radiation_Model/source/OpticallyThick/OpticallyThickIon.cpp',
'../../Radiation_Model/source/OpticallyThick/RadiativeRates.cpp',
'../../Resources/source/gammabeta.cpp',
'../../Resources/source/fitpoly.cpp',
'../../Resources/Utils/generatePieceWiseFit/source/piecewisefit.cpp',
'../../Resources/Utils/regPoly/regpoly.cpp',
'../../Resources/Utils/regPoly/nrutil.cpp',
'../../Resources/source/file.cpp',
]
flags = ['-Wall',] + self.optimization_flags
if self.config['general'].get('use_openmp', False):
flags += ['-fopenmp']
return self.env.get_template('build_script.bat').render(
compiler=self.compiler,
files=files,
flags=flags,
executable='../../HYDRAD.exe',
)