Fix beam heating and OpenMP options (#164)

* make gitpython a core dependency * simplify run_shell_command to just take a path * fix openmp option by adding chunk_size param * The build scripts are now configured on the fly * add a few smoke tests * fiddle with openmp test * mark openmp test as xfail on macos * mark openmp test as xfail on macos * fix parameter table in docs * actually fix the examples * forgot to commit example changes * fix CI badge in readme
rice-solar-physics · Nov 9, 2023 · 391ac2c · 391ac2c
1 parent 52a9666
commit 391ac2c
Show file tree

Hide file tree

Showing 16 changed files with 174 additions and 73 deletions.
diff --git a/README.md b/README.md
@@ -1,6 +1,6 @@
 # pydrad
 
-![pydrad CI status](https://github.com/rice-solar-physics/pydrad/workflows/Test/badge.svg)
+[![pydrad CI status](https://github.com/rice-solar-physics/pydrad/workflows/test.yml/badge.svg?branch=main)](https://github.com/rice-solar-physics/pydrad/actions)
 [![Documentation Status](https://readthedocs.org/projects/pydrad/badge/?version=latest)](https://pydrad.readthedocs.io/en/latest/?badge=latest)
 [![codecov](https://codecov.io/gh/rice-solar-physics/pydrad/branch/master/graph/badge.svg)](https://codecov.io/gh/rice-solar-physics/pydrad)
 

diff --git a/conftest.py b/conftest.py
@@ -126,8 +126,5 @@ def hydrad(tmpdir_factory, hydrad_clean):
     # containing the results.
     hydrad_tmp = tmpdir_factory.mktemp('hydrad_tmp')
     configuration.setup_simulation(hydrad_tmp, hydrad_clean, overwrite=True)
-    pydrad.configure.util.run_shell_command(
-        ['./HYDRAD.exe'],
-        hydrad_tmp,
-    )
+    pydrad.configure.util.run_shell_command(hydrad_tmp / 'HYDRAD.exe')
     return hydrad_tmp
diff --git a/docs/conf.py b/docs/conf.py
@@ -124,7 +124,7 @@ def render_rst_table(section):
     in the config_tables.yml file. Changes to this file will not persist
     after each docs build.
 
-.. _configuration tables:
+.. _configuration-tables:
 
 HYDRAD Configuration Parameters
 ================================

diff --git a/docs/config_tables.yml b/docs/config_tables.yml
@@ -20,6 +20,7 @@ General:
   - minimum_collisional_coupling_timescale
   - force_single_fluid
   - use_openmp
+  - grid_cells_per_thread
   - open_field
   - force_symmetry
   Description:
@@ -47,6 +48,7 @@ General:
   - If true, force electron and ion quantities to be equal
   - If true, parallelize over threads with `OpenMP <https://www.openmp.org/>`__. This
     option is most useful when including a NLTE chromosphere.
+  - If using OpenMP parallelization, approximate number of grid cells assigned to each thread
   - If true, one footpoint is assumed to not connect to the surface
   - ''
   Type:
@@ -70,6 +72,7 @@ General:
   - '``float``'
   - '``bool``'
   - '``bool``'
+  - '``int``'
   - '``bool``'
   - '``bool``'
   Units:
@@ -95,6 +98,7 @@ General:
   - ''
   - ''
   - ''
+  - ''
 Polynomial Fit:
   Name:
   - x

diff --git a/examples/configure_simulation.py b/examples/configure_simulation.py
@@ -6,7 +6,7 @@
 This example demonstrates the various ways to configure a HYDRAD
 simulation.
 """
-import os
+import pathlib
 import tempfile
 
 import astropy.units as u
@@ -27,7 +27,7 @@
 # Mm loop lasting 5000 s heated by a single 200 s nanoflare
 # solved on an adaptive grid.
 # A complete list of configuration parameters can be found in
-# the `configuration tables`_ page.
+# the `configuration-tables`_ page.
 config_dict = {
     'general': {
         'loop_length': 80*u.Mm,
@@ -139,8 +139,8 @@
 # copy of HYDRAD. You can use a copy you already have locally or
 # use the following convenience function to grab the most recent
 # version from GitHub
-tmpdir = tempfile.mkdtemp()  # Change this to wherever you want to save your clean HYDRAD copy
-hydrad_clean = os.path.join(tmpdir, 'hydrad-clean')
+tmpdir = pathlib.Path(tempfile.mkdtemp())  # Change this to wherever you want to save your clean HYDRAD copy
+hydrad_clean = tmpdir / 'hydrad-clean'
 get_clean_hydrad(hydrad_clean, from_github=True)
 
 #################################################################
@@ -150,7 +150,7 @@
 # to our new simulation. This function will write all of the needed
 # aforementioned configuration files and run the hydrostatic solver which
 # will provide the initial conditions for the actual simulation.
-c.setup_simulation(os.path.join(tmpdir, 'test-run'), hydrad_clean)
+c.setup_simulation(tmpdir / 'test-run', hydrad_clean)
 
 #################################################################
 # To avoid having to repeatedly setup the configuration dictionary,
@@ -160,7 +160,7 @@
 # like the config directory and can be easily read and written.
 #
 # To save the configuration to disk and then load it back into a `dict`,
-asdf_config = os.path.join(tmpdir, 'test_config.asdf')
+asdf_config = tmpdir / 'test_config.asdf'
 c.save_config(asdf_config)
 config_from_disk = Configure.load_config(asdf_config)
 print(config_from_disk)

diff --git a/examples/parse_simulation.py b/examples/parse_simulation.py
@@ -7,7 +7,7 @@
 read in and visualize the outputs as a function of time and field-aligned
 spatial coordinate.
 """
-import os
+import pathlib
 import tempfile
 
 import astropy.units as u
@@ -18,7 +18,7 @@
 from pydrad.configure.util import get_clean_hydrad, run_shell_command
 from pydrad.parse import Strand
 
-tmpdir = tempfile.mkdtemp()  # Change to wherever you want to save your clean HYDRAD copy
+tmpdir = pathlib.Path(tempfile.mkdtemp())  # Change to wherever you want to save your clean HYDRAD copy
 
 
 #################################################################
@@ -38,7 +38,7 @@
 #
 # First, we'll setup a simple HYDRAD run and run the simulation.
 # Grab a new copy of HYDRAD
-hydrad_clean = os.path.join(tmpdir, 'hydrad-clean')
+hydrad_clean = tmpdir / 'hydrad-clean'
 get_clean_hydrad(hydrad_clean, from_github=True)
 
 #################################################################
@@ -71,9 +71,9 @@
 # directly from Python, but this is easily done via the command
 # line as well. This can take a few minutes.
 c = Configure(config)
-hydrad_results = os.path.join(tmpdir, 'steady-run')
+hydrad_results = tmpdir / 'steady-run'
 c.setup_simulation(hydrad_results, hydrad_clean)
-run_shell_command(['./HYDRAD.exe'], hydrad_results)
+run_shell_command(hydrad_results / 'HYDRAD.exe')
 
 #################################################################
 # To parse the results of a simulation, we create a `~pydard.parse.Strand`

diff --git a/pydrad/configure/configure.py b/pydrad/configure/configure.py
@@ -30,8 +30,6 @@ class Configure(object):
     ----------
     config : `dict`
         All input parameters for configuring simulation
-    templates : `dict`
-        Templates to override defaults, optional
     """
 
     def __init__(self, config, **kwargs):
@@ -48,6 +46,9 @@ def __init__(self, config, **kwargs):
         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):
@@ -123,17 +124,14 @@ def setup_initial_conditions(self, root_dir, execute=True):
             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),
+            ('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
@@ -150,17 +148,11 @@ def setup_initial_conditions(self, root_dir, execute=True):
             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 / 'Initial_Conditions' / 'build_scripts' / 'build_initial_conditions.bat',
-            mode=stat.S_IRWXU,
-        )
-        run_shell_command(
-            ['./build_initial_conditions.bat'],
-            root_dir / 'Initial_Conditions' / 'build_scripts',
-        )
+        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(['./Initial_Conditions.exe'], root_dir)
+            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)
 
@@ -174,23 +166,17 @@ def setup_hydrad(self, root_dir):
             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),
+            ('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)]
@@ -202,14 +188,9 @@ def setup_hydrad(self, root_dir):
         for filename, filestring in files:
             with (root_dir / filename).open(mode='w') as f:
                 f.write(filestring)
-        # NOTE: using OpenMP requires an alternate compile script
-        if self.config['general'].get('use_openmp', False):
-            build_script = 'build_HYDRAD_OPENMP.bat'
-        else:
-            build_script = 'build_HYDRAD.bat'
         # NOTE: make sure we have needed permissions to run compile script
-        os.chmod(root_dir / 'HYDRAD' / 'build_scripts' / build_script, mode=stat.S_IRWXU)
-        run_shell_command([f'./{build_script}'], root_dir / 'HYDRAD' / 'build_scripts')
+        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
@@ -445,3 +426,69 @@ def maximum_cells(self):
                 {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',
+        )
diff --git a/pydrad/configure/templates/build_script.bat b/pydrad/configure/templates/build_script.bat
@@ -0,0 +1 @@
+{{ compiler }} {{ flags | join(' ') }} {{ files | join(' ') }} -o {{ executable }}
diff --git a/pydrad/configure/templates/hydrad.config.h b/pydrad/configure/templates/hydrad.config.h
@@ -36,7 +36,10 @@
 // **** End of Physics ****
 
 // **** Solver ****
-{% if general.use_openmp %}#define OPENMP{% endif %}
+{% if general.use_openmp -%}
+#define OPENMP
+#define CHUNK_SIZE {{ general.grid_cells_per_thread }}
+{%- endif %}
 #define SAFETY_RADIATION {{ solver.safety_radiation | is_required }}
 #define SAFETY_CONDUCTION {{ solver.safety_conduction | is_required }}
 #define SAFETY_ADVECTION {{ solver.safety_advection | is_required }}

diff --git a/pydrad/configure/templates/initial_conditions.config.h b/pydrad/configure/templates/initial_conditions.config.h
@@ -32,7 +32,7 @@
 #define EPSILON {{ solver.epsilon | is_required }}
 
 // **** Grid ****
-{% if grid.adapt -%}#define ADAPT{%- endif %}
+{% if grid.adapt %}#define ADAPT{% endif %}
 #define MIN_CELLS {{ minimum_cells | is_required }}
 #define MAX_CELLS {{ maximum_cells | is_required }}
 #define MAX_REFINEMENT_LEVEL {{ grid.maximum_refinement_level | is_required }}

diff --git a/pydrad/configure/templates/radiation.config.h b/pydrad/configure/templates/radiation.config.h
@@ -41,3 +41,11 @@
 #define EPSILON_D {{ solver.epsilon_d | is_required }}
 #define EPSILON_R {{ solver.epsilon_r | is_required }}
 // **** End of Solver ****
+
+{% if heating.beam -%}
+// **** Heating ****
+// This needs to be here because the radiation model is needed when constructing the electron
+// beam and it is not imported anywhere else
+#define BEAM_HEATING
+{%- endif %}
+
diff --git a/pydrad/configure/tests/test_templates.py b/pydrad/configure/tests/test_templates.py
@@ -531,6 +531,7 @@ def test_radiation_header(configuration):
 #define EPSILON_D 0.1
 #define EPSILON_R 1.8649415311920072
 // **** End of Solver ****"""
+    assert_ignore_blanks(configuration.radiation_header, header)
 
 
 def test_radiation_config_equilibrium(configuration):