Feature/flash template

docs/user-guide/code-generation/jaffgen.md

@@ -30,7 +30,8 @@ A `jaffgen` run is four phases:
 3. **Load the network** — exactly one `Network`, built from the resolved
    network/label/funcfile/radiation settings.
 4. **Generate** — run [`TemplateParser`](template-syntax.md) on every gathered
-   file and write the result to `--outdir`, keeping the original filename.
+   file and write the result to `--outdir`, preserving each file's path
+   relative to its source directory.
 
 The mental model: _jaffgen is the template engine driven over a file set with
 one loaded network behind it._
@@ -43,13 +44,13 @@ The three input flags are additive, not exclusive:
 
 | Flag         | Adds                                                          |
 | ------------ | ------------------------------------------------------------- |
-| `--indir`    | Every file in a directory (non-recursive)                     |
+| `--indir`    | Every file in a directory (recursive), keeping its structure  |
 | `--files`    | Specific individual files                                     |
 | `--template` | A built-in collection from `jaff/templates/generator/<name>/` |
 
 For `--template`, generator files are collected first; any **preprocessor** file
-(`jaff/templates/preprocessor/<name>/`) whose name does not clash is appended —
-so the generator always wins on a name collision.
+(`jaff/templates/preprocessor/<name>/`) whose relative path does not clash is
+appended — so the generator always wins on a path collision.
 
 ---
 

src/jaff/cli/_jaffgen.py

@@ -20,8 +20,8 @@
 Built-in templates live inside ``jaff/templates/generator/<name>/`` and
 ``jaff/templates/preprocessor/<name>/``.  When ``--template <name>`` is
 given, all files from the *generator* subdirectory are collected first; any
-preprocessor file whose name does not clash with a generator file is appended
-afterwards, so the generator always wins on name collisions.
+preprocessor file whose relative path does not clash with a generator file is
+appended afterwards, so the generator always wins on path collisions.
 
 Usage
 -----
@@ -47,6 +47,7 @@
 from __future__ import annotations
 
 import argparse
+from dataclasses import dataclass
 from inspect import signature
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
@@ -55,9 +56,9 @@
 
 from .. import Network
 from ..cli import ConfigTable
-from ..config import JAFF_DIR, NETWORK_DIR, TEMPLATES_DIR
 from ..codegen import Language, TemplateParser
 from ..common import motd
+from ..config import JAFF_DIR, NETWORK_DIR, TEMPLATES_DIR
 from ..drivers import HDF5, Toml
 from ..io import JaffLogger, jaff_progress
 from ..types import HDF5Dict
@@ -137,7 +138,7 @@ def __init__(self):
         self.network_dir: Path = NETWORK_DIR
         self.generator_template_dir: Path = TEMPLATES_DIR / "generator"
         self.preprocessor_template_dir: Path = TEMPLATES_DIR / "preprocessor"
-        self.files: list[Path] = []
+        self.files: list[FileStruct] = []
         self.jaffgen_config: JaffgenProps = {"netprops": {}}  # type: ignore
         self.jaffgen_config_raw: Toml | None = None
 
@@ -293,7 +294,8 @@ def __process_files(self) -> None:
 
         Iterates over ``self.files``, instantiates a
         :class:`~jaff.codegen.TemplateParser` for each, and writes the
-        generated text to ``output_dir / <filename>``.
+        generated text to ``output_dir / <relative_path>``, recreating any
+        source subdirectories under the output directory.
 
         Returns
         -------
@@ -303,20 +305,20 @@ def __process_files(self) -> None:
         for file in jaff_progress.track(self.files, description="Processing files"):
             # Instantiate the parser for this single template.
             fparser: TemplateParser = TemplateParser(
-                self.net, file, self.jaffgen_config["default_lang"]
+                self.net, file.abspath, self.jaffgen_config["default_lang"]
             )
 
             # Generate the full output text for this file.
             lines: str = fparser.parse_file()
 
             # Write the generated code into the output directory, preserving
-            # the original filename.
-            outfile: Path = self.jaffgen_config["output_dir"] / file.name
-            with open(outfile, "w") as f:
-                f.write(lines)
+            # the file's path relative to its source directory.
+            outfile: Path = self.jaffgen_config["output_dir"] / file.relpath
+            outfile.parent.mkdir(parents=True, exist_ok=True)
+            outfile.write_text(lines)
 
             self.logger.info(
-                f"[cyan]{file.name}[/] created at {self.jaffgen_config['output_dir']}"
+                f"[cyan]{file.relpath.name}[/] created at {self.jaffgen_config['output_dir']}"
             )
 
         self.logger.info("[green]Successfully generated files[/]")
@@ -343,20 +345,20 @@ def __read_jaff_config_from_files(self) -> None:
 
         # Search for a jaff.toml among the already-collected template files.
         jaff_config_index: int | None = next(
-            (i for i, f in enumerate(self.files) if f.name == "jaff.toml"), None
+            (i for i, f in enumerate(self.files) if f.abspath.name == "jaff.toml"), None
         )
 
         if jaff_config_index is not None:
-            self.__set_config(self.files[jaff_config_index])
+            self.__set_config(self.files[jaff_config_index].abspath)
 
     def __set_template(self, template: str | None) -> None:
         """
         Resolve a named built-in template and collect its files.
 
         Looks up *template* inside ``jaff/templates/generator/``.  Generator
-        files are always preferred; any preprocessor file whose name does not
-        match an existing generator file is appended to ``self.files`` as a
-        fallback.
+        files are always preferred; any preprocessor file whose relative path
+        does not match an existing generator file is appended to ``self.files``
+        as a fallback.
 
         Parameters
         ----------
@@ -389,18 +391,22 @@ def __set_template(self, template: str | None) -> None:
         generator_template_path: Path = self.generator_template_dir / template
         preprocessor_template_path: Path = self.preprocessor_template_dir / template
         generator_files = [
-            file for file in generator_template_path.rglob("*") if not file.is_dir()
+            FileStruct(f, f.relative_to(generator_template_path))
+            for f in generator_template_path.rglob("*")
+            if not f.is_dir()
         ]
-        preprocesor_files = [
-            file for file in preprocessor_template_path.rglob("*") if not file.is_dir()
+        preprocessor_files = [
+            FileStruct(f, f.relative_to(preprocessor_template_path))
+            for f in preprocessor_template_path.rglob("*")
+            if not f.is_dir()
         ]
         self.files.extend(generator_files)
 
         # Add preprocessor files only when there is no generator file with the
-        # same name — generator files take precedence.
-        generator_file_names = [file.name for file in generator_files]
-        for file in preprocesor_files:
-            if file.name not in generator_file_names:
+        # same relative path — generator files take precedence.
+        generator_files_relative = [f.relpath for f in generator_files]
+        for file in preprocessor_files:
+            if file.relpath not in generator_files_relative:
                 self.files.append(file)
 
         self.jaffgen_config["template"] = template
@@ -519,7 +525,7 @@ def __set_input_files(self, input_files: list[str] | None) -> None:
             if not infile.is_file():
                 raise FileNotFoundError(f"{file} is not a file")
 
-            self.files.append(infile)
+            self.files.append(FileStruct(infile, Path(infile.name)))
             self.jaffgen_config["input_files"].append(infile)
 
     def __set_input_dir(self, input_dir: str | None) -> None:
@@ -551,8 +557,11 @@ def __set_input_dir(self, input_dir: str | None) -> None:
 
         indir = indir.resolve()
 
-        # Add all regular files in the directory (non-recursive).
-        self.files.extend([f for f in indir.iterdir() if f.is_file()])
+        # Add all regular files in the directory (recursive), keeping each
+        # file's path relative to indir so the structure is preserved.
+        self.files.extend(
+            [FileStruct(f, f.relative_to(indir)) for f in indir.rglob("*") if f.is_file()]
+        )
         self.jaffgen_config["input_dir"] = indir
 
     def __set_funcfile(self, funcfile: str | None) -> None:
@@ -629,8 +638,7 @@ def __set_output_dir(self, output_dir: str | None) -> None:
         outdir = outdir.resolve()
 
         # Create the directory on first use if it does not yet exist.
-        if not outdir.exists():
-            outdir.mkdir()
+        outdir.mkdir(parents=True, exist_ok=True)
 
         if not outdir.is_dir():
             raise NotADirectoryError(f"Output path is not a directory: {outdir}")
@@ -844,6 +852,12 @@ def __handle_data_tables(self, props: list) -> None:
                 )
 
 
+@dataclass(frozen=True)
+class FileStruct:
+    abspath: Path
+    relpath: Path
+
+
 def main():
     """Entry point registered as the ``jaffgen`` console script."""
     JaffGen()

src/jaff/codegen/_template_engine.py

@@ -828,10 +828,16 @@ def __do_iterative_repeat(
             # extras format: [KEY1, VALUE1, KEY2, VALUE2, ...]
             # Step through by 2s to get keys, then check if next value is "TRUE"
             for i, extra in enumerate(extras[::2]):
+                raw_value = extras[2 * i + 1]
+                # Interpret literals (True/False/ints) via literal_eval, but fall
+                # back to the raw string for bare identifiers
+                try:
+                    value = ast.literal_eval(raw_value)
+                except (ValueError, SyntaxError):
+                    value = raw_value
                 # Call the kwargs generator for this extra modifier
-                # extras[2*i+1] gives the corresponding value for this key
                 additional_kwargs = self.__get_special_var_dict[extra]["kwargs"](
-                    extra, ast.literal_eval(extras[2 * i + 1])
+                    extra, value
                 )
                 kwargs = {**kwargs, **additional_kwargs}
 
@@ -1368,7 +1374,9 @@ def __get_parser_dict(self) -> dict[str, CommandProps]:
                     },
                     # Returns: list[str] - species names with +/-
                     "species_with_normalized_sign": {
-                        "func": self.net.species.normalized_names,
+                        "func": lambda **kwargs: self.net.species.normalized_names(
+                            **kwargs
+                        ),
                         "vars": ["idx", "specie_with_normalized_sign"],
                     },
                     # Returns: list[str] - element symbols
@@ -1679,6 +1687,8 @@ def __get_special_var_dict(self) -> dict[str, dict[str, Any]]:
             "RAD_ORDER": {"kwargs": lambda var, value: {"rad_order": value}},
             "SPECIFIC_EINT": {"kwargs": lambda var, value: {"specific_eint": value}},
             "NORM": {"kwargs": lambda var, value: {"norm": value}},
+            "POS": {"kwargs": lambda var, value: {"pos": value}},
+            "NEG": {"kwargs": lambda var, value: {"neg": value}},
         }
 
         return svar_dict

src/jaff/templates/generator/flash/KromeChemistryMain/Config

@@ -0,0 +1,20 @@
+# JAFF-generated FLASH Config for chemistry network.
+# Auto-generated by jaffgen (template collection: flash).
+
+# The species composition lives on the Simulation side, keyed by the network
+# label. $label$ expands to the jaffgen network label (e.g. h_photo).
+# $JAFF SUB label
+REQUIRES Simulation/SimulationComposition/KromeChemistry/$label$
+# $JAFF END
+
+# Network variables consumed by the generated rate expressions.
+PARAMETER krome_crate REAL 1.0e-17
+PARAMETER krome_av    REAL 0.0
+PARAMETER krome_temp_floor        REAL 1.0
+PARAMETER krome_density_threshold REAL 0.0
+
+# One SPECIES per network species, emitted in jaffgen order so the FLASH
+# unk indices line up with the idx_* parameters in the commons module.
+# $JAFF REPEAT idx, specie_with_normalized_sign IN species_with_normalized_sign $[POS j NEG k]$
+SPECIES $specie_with_normalized_sign$  # $idx$
+# $JAFF END

src/jaff/templates/generator/flash/KromeChemistryMain/KromeChemistry.F90

@@ -0,0 +1,128 @@
+!!****if* source/physics/sourceTerms/KromeChemistry/KromeChemistry
+!!
+!! NAME
+!!  KromeChemistry
+!!
+!! SYNOPSIS
+!!  call KromeChemistry(integer :: blockCount,
+!!                      integer, dimension(blockCount) :: blockList,
+!!                      real :: dt)
+!!
+!! DESCRIPTION
+!!  Apply the JAFF-generated chemistry network to all cells in the given blocks.
+!!  Auto-generated by jaffgen (template collection: flash).
+!!
+!!  The network right-hand side (fex) comes from the generated ode module; the
+!!  ODE integrator is FLASH's own ODEPACK/DLSODES (opkda1/opkda2/opkdmain),
+!!  shared by the KromeChemistry unit. jaffgen supplies no integrator.
+!!
+!!***
+
+subroutine KromeChemistry(blockCount, blockList, dt)
+
+  use KromeChemistry_data
+  use commons, ONLY : nspecs, idx_tgas
+  use ode, ONLY : fex, jes
+  use Timers_interface, ONLY : Timers_start, Timers_stop
+  use Grid_interface, ONLY : Grid_getBlkIndexLimits, Grid_getBlkPtr, Grid_releaseBlkPtr
+  use Eos_interface, ONLY : Eos_wrapped
+
+  implicit none
+
+#include "constants.h"
+#include "Flash.h"
+#include "Eos.h"
+
+  integer, intent(in) :: blockCount
+  integer, intent(in), dimension(blockCount) :: blockList
+  real, intent(in) :: dt
+
+  integer :: i, j, k, n, blockID, thisBlock
+  integer, dimension(2,MDIM) :: blkLimits, blkLimitsGC
+  real, dimension(NSPECIES) :: x, xn
+  real :: tmp, rho
+
+  ! DLSODES work space. neq = species + one gas-temperature slot.
+  integer, parameter :: neq = NSPECIES + 1
+  integer, parameter :: lrw = 20 + 3*neq**2 + 20*neq
+  integer, parameter :: liw = 30
+  real    :: rwork(lrw), rtol(neq), atol(neq), nvec(neq), tloc, tend
+  integer :: iwork(liw), itol, itask, istate, iopt, mf
+
+  real, pointer, dimension(:,:,:,:) :: solnData
+
+  if (.not. useKromeChemistry) return
+
+  call Timers_start("chemistry")
+
+  ! Integrator controls (see the DLSODES documentation in opkdmain.f).
+  itol  = 4      ! rtol and atol are both arrays
+  itask = 1      ! normal integration to tend
+  iopt  = 0      ! no optional inputs
+  mf    = 222    ! stiff, internally generated sparse Jacobian
+  rtol  = 1.0e-6
+  atol  = 1.0e-6
+
+  do thisBlock = 1, blockCount
+
+     blockID = blockList(thisBlock)
+
+     call Grid_getBlkIndexLimits(blockID, blkLimits, blkLimitsGC)
+     call Grid_getBlkPtr(blockID, solnData)
+
+     do k = blkLimits(LOW,KAXIS), blkLimits(HIGH,KAXIS)
+        do j = blkLimits(LOW,JAXIS), blkLimits(HIGH,JAXIS)
+           do i = blkLimits(LOW,IAXIS), blkLimits(HIGH,IAXIS)
+
+              tmp = solnData(TEMP_VAR,i,j,k)
+              rho = solnData(DENS_VAR,i,j,k)
+
+              ! skip cells below the chemistry density threshold
+              if (rho .lt. krome_density_threshold) cycle
+
+              ! mass fractions -> number densities
+              do n = 1, NSPECIES
+                 x(n)  = solnData(specieMap(n),i,j,k)
+                 xn(n) = x(n) * rho / mp / amu(n)
+              enddo
+
+              ! assemble the DLSODES state vector: species then temperature
+              nvec(1:NSPECIES) = xn(:)
+              nvec(idx_tgas)   = max(tmp, krome_temp_floor)
+
+              tloc   = 0.0
+              tend   = dt
+              istate = 1
+              rwork  = 0.0
+              iwork  = 0
+
+              call DLSODES(fex, neq, nvec, tloc, tend, itol, rtol, atol, &
+                           itask, istate, iopt, rwork, lrw, iwork, liw, jes, mf)
+
+              ! unpack integrated state
+              xn(:) = nvec(1:NSPECIES)
+              tmp   = nvec(idx_tgas)
+
+              ! number densities -> mass fractions
+              do n = 1, NSPECIES
+                 x(n) = xn(n) * mp * amu(n) / rho
+                 solnData(specieMap(n),i,j,k) = x(n)
+              enddo
+
+              solnData(TEMP_VAR,i,j,k) = tmp
+
+           enddo
+        enddo
+     enddo
+
+     call Grid_releaseBlkPtr(blockID, solnData)
+
+     call Eos_wrapped(MODE_DENS_TEMP, blkLimits, blockID)
+
+  enddo
+
+  call Timers_stop("chemistry")
+
+  return
+
+end subroutine KromeChemistry